diff --git a/fern/versions/latest/pages/training-tutorials/external-agent-harnesses.mdx b/fern/versions/latest/pages/training-tutorials/external-agent-harnesses.mdx index 2c3d5de6dd..697a1ddb99 100644 --- a/fern/versions/latest/pages/training-tutorials/external-agent-harnesses.mdx +++ b/fern/versions/latest/pages/training-tutorials/external-agent-harnesses.mdx @@ -70,7 +70,11 @@ This failure is silent. Check `n_calls` on the first run before relying on the r ## What you get back -Each rollout's model calls are stitched into a single Responses payload with contiguous `output` items. Each item's `prompt_token_ids` contains the running sequence. Its `generation_token_ids` contains the tokens sampled by the policy at that step. Gym replaces the rollout's `response.output` with these items, so a trainer reads `response.output` the same way for native agents and external harnesses. +Each captured call records one parent-resolution result. `ROOT` means that the request contains no previous model-authored output. `RESOLVED` means that exactly one committed call matches and its request context was verified. `UNRESOLVED` means that the request appears to continue prior model output, but Gym cannot prove which committed call produced it. + +Each rollout's resolved model calls are stitched into a single Responses payload with contiguous `output` items. Each item's `prompt_token_ids` contains the running sequence. Its `generation_token_ids` contains the tokens sampled by the policy at that step. Gym replaces the rollout's `response.output` with these items, so a trainer reads `response.output` the same way for native agents and external harnesses. + +An unresolved call begins a separate incomplete fragment. Gym never crosses that boundary with token-prefix inference. The call and its descendants remain available for diagnostics or a consumer that explicitly supports partial trajectories, but single-response delivery sets `mask_sample: true`. Prefix inference is reserved for records written before parent-resolution metadata existed. The loss mask follows from this structure instead of being sent separately. Prompt positions provide context. Generation positions are trainable. @@ -87,7 +91,8 @@ Gym attaches a metrics dictionary to each rollout under `_ng_token_capture`. Agg | `delivered_fraction` | 1.0 | Sampled tokens were captured but not delivered. | | `quarantined_calls` | 0 | Two calls could not be told apart, so neither was used. | | `empty_generation_calls` | 0 | The output budget or a content filter is truncating generations. | -| `mask_sample` | absent | The rollout lost a call and must not be trained on. | +| `unresolved_parent_calls` | 0 | A request appeared to continue prior model output, but no exact committed parent was proven. | +| `mask_sample` | absent | The rollout lost a call, split at an unresolved parent, or otherwise cannot be trained safely. | Pay particular attention to `n_calls`. A value of exactly 1 means the agentic path was never exercised, even though every other key can look correct. @@ -99,7 +104,7 @@ Gym defines the record shape and builds each record. The training framework cont ### The interfaces -Gym describes the sink and source as structural protocols. Framework adapters do not inherit from these definitions or import them at runtime. They implement the same method signatures, and Gym consumes the resulting objects by that method shape. The definitions below are the reference contract. +Gym describes the sink, source, and lineage resolver as structural protocols. Framework adapters do not inherit from these definitions or import them at runtime. They implement the same method signatures, and Gym consumes the resulting objects by that method shape. The definitions below are the reference contract. ```python class TokenSink(Protocol): @@ -111,9 +116,14 @@ class TokenSource(Protocol): async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: ... async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: ... async def close(self) -> None: ... + +class LineageStore(Protocol): + async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: ... + def is_process_shared(self) -> bool: ... + async def close(self) -> None: ... ``` -`put` must make the record durable before it returns. A reader that runs after the rollout must see every acknowledged record. This guarantee allows the consumer to freeze one complete view after the harness finishes. +`TokenSink` remains the writer-side interface. `put` is the only publication boundary: it must make the complete token record and its compact continuation lookup metadata durable before returning. `LineageStore` is the worker-side read client over those committed entries; it does not publish a second record. A later request may resolve an entry only after the corresponding `put` has returned. `mark_incomplete` is the durable signal that a rollout lost a call. The model call still succeeds. A sink that drops this signal makes an incomplete rollout look complete. @@ -123,9 +133,11 @@ class TokenSource(Protocol): `close` releases client resources. Gym closes clients it constructs, but it does not close a caller-installed source. +`LineageResolution` has three semantic outcomes: `ROOT`, `RESOLVED`, and `UNRESOLVED`. Diagnostic reasons such as an ambiguous lookup or unavailable backend may accompany `UNRESOLVED`, but they do not change reconstruction behavior. + ### Connecting a framework-owned transport -The training framework owns the transport and its client configuration. Gym owns the model-server processes, so each server worker constructs a framework-provided `TokenSink` proxy from the configured class path. The configured class is a worker factory descriptor, not a shared Python object or a transport implementation owned by Gym. +The training framework owns the transport and its client configuration. Gym owns the model-server processes, so each server worker constructs framework-provided `TokenSink` and `LineageStore` proxies from configured class paths. The configured classes are worker factory descriptors, not shared Python objects or transport implementations owned by Gym. ```yaml env: @@ -136,10 +148,16 @@ env: sink_kwargs: endpoint: ${oc.env:MY_DATAPLANE_URL} shard: ${oc.select:cluster_shard,0} + lineage_store: my_pkg.sinks:MyDataPlaneLineageStore + lineage_store_kwargs: + endpoint: ${oc.env:MY_DATAPLANE_URL} + shard: ${oc.select:cluster_shard,0} rebuild_response: false ``` -Gym passes `sink_kwargs` to the constructor, so a sink can receive the required endpoint, client, or credentials instead of reading ambient state. Use `${oc.env:VAR}` for secrets instead of writing them into the config. Unsupported constructor arguments cause a startup error. A sink that does not implement `mark_incomplete` also causes a startup error because it could otherwise make a rollout with a missing call look complete. +Gym passes each kwargs block to its constructor, so the clients can receive the required endpoint, shard, or credentials instead of reading ambient state. Use `${oc.env:VAR}` for secrets instead of writing them into the config. Unsupported constructor arguments cause a startup error. A sink that does not implement `mark_incomplete` also causes a startup error because it could otherwise make a rollout with a missing call look complete. + +The sink and lineage resolver must use the same backend namespace. The backend transaction behind `TokenSink.put` stores the `TokenEntry` and updates the continuation-key index together. Pointing the two clients at unrelated services does not satisfy the protocol even if both methods return successfully. `sink` replaces the file store, so a `dir` configured alongside it is not used. This condition produces a warning instead of an error because no data is lost. No capture files appear on disk. @@ -163,9 +181,11 @@ Configure the sink instead of installing it from a launcher script. Programmatic `install_token_sink` sets a process global. A model server with `num_workers > 1` launches uvicorn with an app string and `workers=N`. Uvicorn spawns workers that re-import the app module instead of inheriting the launcher's memory. A sink installed by the parent process therefore does not exist in any worker. Capture then falls back to the file store. If no `dir` is set, the worker has no local destination. -Gym constructs the configured sink inside each worker at app startup, which avoids this process-boundary problem. +Gym constructs configured sink and resolver clients inside each worker at app startup. Every client must connect to the same process-shared backend. `LineageStore.is_process_shared()` must return `true`; startup rejects a process-local resolver when `num_workers > 1`. +With Gym's local file backend, all workers append and resolve from the same token JSONL under one per-rollout lock. The resolver incrementally indexes only newly appended entries in each worker. Freeze and retirement use the same lock. Retirement keeps the lock file and a tombstone, so an old worker cannot recreate a retired capture. + ### Reading records back `TokenCaptureStore` implements both protocols and is the default. A reader beside the store uses the store as its `TokenSource`. This arrangement applies to `gym eval run` and to a trainer colocated with the model server. The store directory should therefore be node-local. diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 5d323dad33..d9a2ac140c 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -71,15 +71,18 @@ from nemo_gym.token_id_capture import ( CaptureContext, capture_tokens, + installed_lineage_store, installed_token_sink, register_call_intent, reset_token_sink, + resolve_parent, set_token_sink, ) # The store factory needs Gym's server stack. # The leaf package does not re-export it. from nemo_gym.token_id_capture.config import token_id_capture_config +from nemo_gym.token_id_capture.lineage import FileLineageStore from nemo_gym.token_id_capture.store import make_token_store @@ -90,6 +93,65 @@ _ANTHROPIC_CONVERTER = AnthropicConverter() +def _request_messages(body: Any) -> list[dict]: + """Return the conversation carried by any supported dialect. + + Lineage uses model-authored turns to identify the parent call. + Chat and Anthropic use ``messages``. + Responses uses ``input``. + The request envelope is prepended as a pseudo-turn because it also shapes the prompt. + """ + if body is None: + return [] + getter = body.get if isinstance(body, dict) else lambda key, default=None: getattr(body, key, default) + messages = getter("messages", None) + if isinstance(messages, list): + turns = [m if isinstance(m, dict) else m.model_dump() for m in messages if m is not None] + else: + items = getter("input", None) + turns = ( + [i if isinstance(i, dict) else i.model_dump() for i in items if i is not None] + if isinstance(items, list) + else [] + ) + return _request_envelope(getter) + turns + + +# This role cannot collide with a dialect role. +# It is not assistant-authored and does not affect the lookup fingerprint. +_ENVELOPE_ROLE = "_ng_request_envelope" + + +def _request_envelope(getter: Any) -> list[dict]: + """Return prompt-shaping request fields that are not turns. + + Instructions and tools can be siblings of the message list. + The chat template renders both into the prompt. + Including them prevents prefix reuse across different request envelopes. + """ + instructions = getter("instructions", None) + tools = getter("tools", None) + if not instructions and not tools: + return [] + envelope = {"instructions": _plain(instructions), "tools": _plain(tools)} + return [{"role": _ENVELOPE_ROLE, "content": json.dumps(envelope, sort_keys=True, default=str)}] + + +def _plain(value: Any) -> Any: + """Reduce a request field to plain data so equal schemas serialize equally. + + Handlers can expose tools as dictionaries or Pydantic models. + Normalization keeps an unchanged schema stable across handlers. + """ + if hasattr(value, "model_dump"): + return value.model_dump() + if isinstance(value, (list, tuple)): + return [_plain(item) for item in value] + if isinstance(value, dict): + return {key: _plain(item) for key, item in value.items()} + return value + + class BaseResponsesAPIModelConfig(BaseRunServerInstanceConfig): pass @@ -109,6 +171,7 @@ def setup_webserver(self) -> FastAPI: capture_config, model_server_name=self.config.name, global_config_dict=self.server_client.global_config_dict, + num_workers=self.config.num_workers, ) app.post("/v1/chat/completions")(self.chat_completions_dispatch) @@ -210,12 +273,19 @@ async def _invoke_chat_completions( # chat_completions() signatures vary across servers: some take a leading `request`, some # only `body`. Dispatch on whichever this server declares so the shared dispatch works for # all of them. + # Resolve the parent from the received request before dispatch. + # Exact prefix supply and capture share this decision. + request_messages = _request_messages(params) + await resolve_parent(request_messages) await register_call_intent() if "request" in inspect.signature(self.chat_completions).parameters: completion = await self.chat_completions(request=request, body=params) else: completion = await self.chat_completions(body=params) - await capture_tokens(completion) + await capture_tokens( + completion, + request_messages=request_messages, + ) return completion async def messages(self, request: Request, body: dict = Body()): @@ -244,6 +314,10 @@ async def _invoke_responses( # responses() signatures vary across servers: some take a leading `request`, some only # `body`. Dispatch on whichever this server declares so the default messages() works for # all of them. + # Resolve the parent from the received request before dispatch. + # Exact prefix supply and capture share this decision. + request_messages = _request_messages(params) + await resolve_parent(request_messages) await register_call_intent() if "request" in inspect.signature(self.responses).parameters: response = await self.responses(request=request, body=params) @@ -252,7 +326,10 @@ async def _invoke_responses( # Capture before streaming dispatch wraps the response. # Anthropic mapping drops the token fields. # The assembled response still carries them here for every dialect. - await capture_tokens(response) + await capture_tokens( + response, + request_messages=request_messages, + ) return response @@ -1068,6 +1145,7 @@ def __init__( model_server_name: str | None, token_store: Any = None, configured_sink: Any = None, + lineage_store: Any = None, token_capture_enabled: bool = False, ) -> None: self._app = app @@ -1077,6 +1155,7 @@ def __init__( self._token_store = token_store # Built from token_id_capture.sink, once, in this process. self._configured_sink = configured_sink + self._lineage_store = lineage_store # Capture may have no destination in this process. # A framework may stage records from its inference worker. # This process still resolves the capture identity. @@ -1138,7 +1217,12 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None sink_token = None if capture_wanted: sink_token = set_token_sink( - CaptureContext(rollout_id=rollout_id, model_call_id=model_call_id, token_sink=token_sink) + CaptureContext( + rollout_id=rollout_id, + model_call_id=model_call_id, + token_sink=token_sink, + lineage_store=self._lineage_store, + ) ) # Training-only capture has no evaluation record. @@ -1308,6 +1392,7 @@ def install_model_call_capture( *, model_server_name: str | None = None, global_config_dict: Any = None, + num_workers: int | None = None, ) -> None: """Install model-call capture middleware. @@ -1319,22 +1404,36 @@ def install_model_call_capture( That path provides a request-scoped token sink. The model server records token ids from its complete response. Consumers access records through ``TokenSource.freeze``. - There is no HTTP token reader. """ + capture_settings = token_id_capture_config(global_config_dict) if global_config_dict is not None else None token_store = make_token_store(global_config_dict) if global_config_dict is not None else None # Build this sink at app startup. # Each uvicorn worker constructs its own sink. # Spawned workers do not inherit a launcher-installed sink. - configured_sink = ( - token_id_capture_config(global_config_dict).build_sink() if global_config_dict is not None else None - ) - owned_sinks = [sink for sink in (configured_sink, token_store) if sink is not None] + configured_sink = capture_settings.build_sink() if capture_settings is not None else None + configured_lineage = capture_settings.build_lineage_store() if capture_settings is not None else None + lineage_store = configured_lineage or installed_lineage_store() + default_lineage = None + if lineage_store is None and token_store is not None: + default_lineage = FileLineageStore(token_store.root) + lineage_store = default_lineage + if capture_settings is not None and capture_settings.enabled and (num_workers or 1) > 1: + if lineage_store is None or not lineage_store.is_process_shared(): + raise ValueError( + "token_id_capture with num_workers > 1 requires a process-shared lineage resolver " + "over the same backend used by the token sink" + ) + owned_endpoints = [ + endpoint + for endpoint in (configured_sink, token_store, configured_lineage, default_lineage) + if endpoint is not None + ] - async def _close_token_sinks() -> None: - for sink in owned_sinks: - await sink.close() + async def _close_capture_endpoints() -> None: + for endpoint in owned_endpoints: + await endpoint.close() - if owned_sinks: + if owned_endpoints: original_lifespan = app.router.lifespan_context @asynccontextmanager @@ -1343,7 +1442,7 @@ async def _capture_lifespan(application): async with original_lifespan(application) as state: yield state finally: - await _close_token_sinks() + await _close_capture_endpoints() app.router.lifespan_context = _capture_lifespan app.add_middleware( @@ -1352,9 +1451,8 @@ async def _capture_lifespan(application): model_server_name=model_server_name, token_store=token_store, configured_sink=configured_sink, - token_capture_enabled=( - token_id_capture_config(global_config_dict).enabled if global_config_dict is not None else False - ), + lineage_store=lineage_store, + token_capture_enabled=capture_settings.enabled if capture_settings is not None else False, ) diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py index de1b3cf97f..99d4f4a906 100644 --- a/nemo_gym/token_id_capture/__init__.py +++ b/nemo_gym/token_id_capture/__init__.py @@ -23,7 +23,6 @@ Its ``snapshot_id`` identifies the exact frozen state. ``TokenCaptureStore`` is Gym's local sink and source implementation. Framework transports may provide their own sink and source. -There is no HTTP token reader. This leaf package avoids imports from Gym's server stack. The rollout-record finalizer needs Gym's server stack. It is deliberately not re-exported here. @@ -39,7 +38,6 @@ from nemo_gym.token_id_capture.builder import ( Chain, assert_prefix_contiguity, - per_request, prefix_merging, project_chain_to_output_items, project_main_chain_response, @@ -52,28 +50,50 @@ trajectories_for_rollout, trajectories_from_source, ) +from nemo_gym.token_id_capture.lineage import ( + FileLineageStore, + IncrementalLineageStore, + InMemoryLineageStore, + LineageIndex, + RolloutLineage, + assistant_fingerprint, + canonicalize_tool_arguments, + stamp_continuation, +) from nemo_gym.token_id_capture.protocols import ( + LineageMatch, + LineageResolution, + LineageStore, TokenCaptureSnapshot, TokenSink, TokenSource, + install_lineage_store, install_token_sink, install_token_source, + installed_lineage_store, installed_token_sink, installed_token_source, ) from nemo_gym.token_id_capture.records import ( + TOKEN_ENTRY_MIN_SCHEMA_VERSION, TOKEN_ENTRY_RECORD_SCHEMA_VERSION, TOKEN_FIELDS, + ParentResolutionStatus, TokenEntry, + compute_digest, + cumulative_tokens, extract_token_fields, + stamp_lineage, ) from nemo_gym.token_id_capture.sink import ( CaptureContext, + capture_health_snapshot, capture_tokens, commit_entry, current_capture_context, register_call_intent, reset_token_sink, + resolve_parent, set_token_sink, ) from nemo_gym.token_id_capture.store import TokenCaptureStore, make_token_store, validate_rollout_id @@ -83,14 +103,24 @@ "Chain", "TokenIdCaptureConfig", "TokenEntry", + "TOKEN_ENTRY_MIN_SCHEMA_VERSION", "TOKEN_ENTRY_RECORD_SCHEMA_VERSION", "TOKEN_FIELDS", + "ParentResolutionStatus", "extract_token_fields", + "compute_digest", + "cumulative_tokens", + "stamp_lineage", "TokenCaptureStore", "validate_rollout_id", "make_token_store", "TokenSink", "TokenSource", + "LineageMatch", + "LineageResolution", + "LineageStore", + "install_lineage_store", + "installed_lineage_store", "TokenCaptureSnapshot", "install_token_sink", "install_token_source", @@ -98,12 +128,20 @@ "installed_token_source", "CaptureContext", "set_token_sink", - "reset_token_sink", + "capture_health_snapshot", "register_call_intent", + "reset_token_sink", + "resolve_parent", "capture_tokens", "commit_entry", "current_capture_context", - "per_request", + "FileLineageStore", + "IncrementalLineageStore", + "InMemoryLineageStore", + "LineageIndex", + "RolloutLineage", + "assistant_fingerprint", + "stamp_continuation", "prefix_merging", "project_chain_to_output_items", "project_main_chain_response", diff --git a/nemo_gym/token_id_capture/builder.py b/nemo_gym/token_id_capture/builder.py index 5a45d58765..277f06c46e 100644 --- a/nemo_gym/token_id_capture/builder.py +++ b/nemo_gym/token_id_capture/builder.py @@ -18,7 +18,6 @@ The builder consumes ``TokenEntry`` records from a ``TokenCaptureSnapshot``. The snapshot is frozen before the consumer passes its entries to the builder. -``per_request`` creates one training sequence per call. It does not infer relationships between calls. It can return multiple trajectories. @@ -45,7 +44,7 @@ from dataclasses import dataclass, field from typing import Callable -from nemo_gym.token_id_capture.records import TokenEntry +from nemo_gym.token_id_capture.records import ParentResolutionStatus, TokenEntry, compute_digest @dataclass @@ -97,6 +96,10 @@ class BuildNotes: unresolved_retries: list[str] = field(default_factory=list) # Calls without generated tokens are excluded from the chain. empty_generation_calls: list[str] = field(default_factory=list) + # Count why recorded parent links were not used. + parent_link_failures: dict[str, int] = field(default_factory=dict) + # These calls begin fragments after an unproven parent boundary. + unresolved_parent_calls: list[str] = field(default_factory=list) @dataclass @@ -106,15 +109,6 @@ class BuildOutput: notes: BuildNotes = field(default_factory=lambda: BuildNotes(builder="")) -def per_request(entries: list[TokenEntry]) -> BuildOutput: - ordered = sorted(entries, key=lambda e: (len(e.prompt_token_ids), e.model_call_id)) - chains = [ - Chain(chain_id=f"req-{i}", root_prompt=list(e.prompt_token_ids), links=[ChainLink(entry=e, interstitial=[])]) - for i, e in enumerate(ordered) - ] - return BuildOutput(chains=chains, notes=BuildNotes(builder="per_request", chains=len(chains))) - - @dataclass(eq=False) # Identity-based equality keeps nodes hashable. class _Node: entry: TokenEntry @@ -122,6 +116,7 @@ class _Node: parent: "_Node | None" = None children: list["_Node"] = field(default_factory=list) quarantined: bool = False + unresolved_boundary: bool = False @dataclass @@ -155,17 +150,81 @@ def infer_parent(self, prompt: list[int]) -> tuple[_Node | None, bool]: return (best[0], len(best) > 1) if best else (None, False) -def _infer_parent(prompt: list[int], index: _PrefixIndex) -> tuple["_Node | None", bool]: - """Infer the parent from the longest cumulative prefix. +def _resolve_parent( + node: "_Node", + by_call_id: dict[str, "_Node"], + prefix_index: _PrefixIndex, +) -> tuple["_Node | None", bool, str | None]: + """Find this call's parent. - This is the fallback when no verified parent link exists. - Identical cumulative sequences are ambiguous. - The caller quarantines an ambiguous subtree. + New records preserve the request-time parent decision. + Token-prefix matching recovers a resolved link whose parent is absent from this build. + This can happen when the parent had an empty generation and was filtered out. + ``note`` reports why the recorded link was not used as recorded. """ - return index.infer_parent(prompt) + prompt = list(node.entry.prompt_token_ids) + resolution = node.entry.parent_resolution + if resolution == ParentResolutionStatus.ROOT: + return None, False, None + if resolution == ParentResolutionStatus.UNRESOLVED: + return None, False, "parent_unresolved" + claimed = node.entry.parent_call_id + if resolution == ParentResolutionStatus.RESOLVED or claimed is not None: + if claimed is None: + return None, False, "resolved_parent_missing_id" + parent = by_call_id.get(claimed) + if parent is None: + # An absent parent is not evidence of conflict. + # Prefix matching may attach the child to a verified surviving ancestor. + # A digest mismatch remains a hard boundary. + inferred, ambiguous = prefix_index.infer_parent(prompt) + if inferred is not None and not ambiguous: + return inferred, False, "parent_call_id_missing_recovered" + return None, False, "parent_call_id_missing" + cum_len = parent.entry.cum_len + if cum_len is None: + cum_len = len(parent.cumulative) + if cum_len <= len(prompt) and compute_digest(prompt[:cum_len]) == ( + parent.entry.digest or compute_digest(parent.cumulative) + ): + return parent, False, None + return None, False, "parent_digest_mismatch" + # Every supported record carries a request-time decision. + # Never guess a parent for a nonconforming record. + return None, False, "missing_resolution" + + +class _NullPrefixIndex: + """Stands in when no entry can need prefix inference. + + Every supported entry carries a request-time parent decision. + Only missing-parent recovery needs the trie. + Avoiding the trie is significant for long multi-call rollouts. + """ + + def add(self, candidate: "_Node") -> None: + return + + def infer_parent(self, prompt: list[int]) -> tuple["_Node | None", bool]: + return None, False def prefix_merging(entries: list[TokenEntry]) -> BuildOutput: + # An at-least-once transport can deliver one entry twice. + # Conflicting payloads for one id are corrupt. + deduped: dict[str, TokenEntry] = {} + duplicate_conflicts: list[str] = [] + for candidate in entries: + previous = deduped.get(candidate.model_call_id) + if previous is None: + deduped[candidate.model_call_id] = candidate + elif (list(previous.prompt_token_ids), list(previous.generation_token_ids)) != ( + list(candidate.prompt_token_ids), + list(candidate.generation_token_ids), + ): + duplicate_conflicts.append(candidate.model_call_id) + entries = list(deduped.values()) + # 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. @@ -175,7 +234,8 @@ def prefix_merging(entries: list[TokenEntry]) -> BuildOutput: entries = [e for e in entries if e.generation_token_ids] if not entries: return BuildOutput( - chains=[], notes=BuildNotes(builder="prefix_merging", empty_generation_calls=empty_generation) + chains=[], + notes=BuildNotes(builder="prefix_merging", empty_generation_calls=empty_generation), ) # Increasing prompt length defines an order from the tokens. @@ -186,12 +246,29 @@ def prefix_merging(entries: list[TokenEntry]) -> BuildOutput: nodes: list[_Node] = [] roots: list[_Node] = [] quarantined: list[str] = [] - prefix_index = _PrefixIndex() - + # The trie exists only for missing-parent recovery. + surviving_ids = {e.model_call_id for e in ordered} + needs_prefix_index = any(e.parent_call_id is not None and e.parent_call_id not in surviving_ids for e in ordered) + prefix_index = _PrefixIndex() if needs_prefix_index else _NullPrefixIndex() + + nodes_by_call_id: dict[str, _Node] = {} + parent_link_failures: dict[str, int] = {} + unresolved_parent_calls: list[str] = [] + for call_id in duplicate_conflicts: + parent_link_failures["duplicate_call_id_conflict"] = ( + parent_link_failures.get("duplicate_call_id_conflict", 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)) - parent, ambiguous = _infer_parent(prompt, prefix_index) + parent, ambiguous, note = _resolve_parent(node, nodes_by_call_id, prefix_index) + if note: + parent_link_failures[note] = parent_link_failures.get(note, 0) + 1 + # A recovered fallback found a safe parent. + if not note.endswith("_recovered"): + node.unresolved_boundary = True + unresolved_parent_calls.append(entry.model_call_id) if parent is not None: node.parent = parent if ambiguous: @@ -203,6 +280,7 @@ def prefix_merging(entries: list[TokenEntry]) -> BuildOutput: roots.append(node) nodes.append(node) prefix_index.add(node) + nodes_by_call_id[entry.model_call_id] = node # Resolve retry siblings. # A harness can retry after a timeout, server error, or dropped stream. @@ -314,12 +392,13 @@ def selection_key(c: Chain) -> tuple: delivered_fraction=round(delivered / captured, 4) if captured else 0.0, unresolved_retries=unresolved_retries, empty_generation_calls=empty_generation, + parent_link_failures=parent_link_failures, + unresolved_parent_calls=unresolved_parent_calls, ) return BuildOutput(chains=chains, quarantined=quarantined, notes=notes) _BUILDERS: dict[str, Callable[[list[TokenEntry]], BuildOutput]] = { - "per_request": per_request, "prefix_merging": prefix_merging, } @@ -391,8 +470,6 @@ def project_main_chain_response(rollout_id: str, out: BuildOutput, model: str = """ if not out.chains: raise ValueError("capture produced no safe trainable chain") - if out.notes.builder == "per_request" and len(out.chains) != 1: - raise ValueError("per_request produced multiple trajectories for a single-response delivery") mains = [c for c in out.chains if c.chain_id == "main"] or out.chains[:1] output = project_chain_to_output_items(mains[0]) if not any(item.get("generation_token_ids") for item in output): diff --git a/nemo_gym/token_id_capture/config.py b/nemo_gym/token_id_capture/config.py index cc0ed87ff3..508a2b233d 100644 --- a/nemo_gym/token_id_capture/config.py +++ b/nemo_gym/token_id_capture/config.py @@ -43,7 +43,6 @@ Consumers construct and inject their ``TokenSource`` in their own process. Consumers call ``TokenSource.freeze`` to obtain an atomic snapshot. Consumers retire that exact snapshot with its ``snapshot_id`` and version. -There is no HTTP token reader. Uvicorn workers use spawned processes. They do not inherit a sink installed by a launcher. Configure the sink here so each worker builds its own. @@ -73,7 +72,9 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator from nemo_gym.token_id_capture.protocols import ( + LineageStore, TokenSink, + installed_lineage_store, installed_token_sink, ) @@ -101,10 +102,18 @@ class TokenIdCaptureSettings(BaseModel): # A real transport needs explicit endpoint, client, or credential wiring. # Use ``${oc.env:VAR}`` for secrets instead of writing them here. sink_kwargs: dict[str, Any] = Field(default_factory=dict) + # Optional process-shared resolver over entries committed by the sink. + # Both clients must use the same backend namespace. + lineage_store: str | None = None + lineage_store_kwargs: dict[str, Any] = Field(default_factory=dict) # Whether Gym freezes capture records and rebuilds the response. # Finalization does not retire the frozen snapshot. # Durable delivery permits retirement by snapshot id and version. rebuild_response: bool = True + # A custom sink normally needs a resolver over the same backend namespace. + # 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. max_mask_fraction: float | None = None @@ -136,12 +145,14 @@ def _validate(self) -> "TokenIdCaptureConfig": "the file store, so %s will not be written to.", block.dir, ) + self._require_resolver(block) return self directory = self.resolved_dir() if directory is None: # A programmatic sink replaces the file store. # That process does not need a directory. if installed_token_sink() is not None: + self._require_resolver(block) return self if not block.rebuild_response: return self @@ -150,6 +161,27 @@ def _validate(self) -> "TokenIdCaptureConfig": raise ValueError("training-token capture directory must be an absolute path") return self + @staticmethod + def _require_resolver(block: TokenIdCaptureSettings) -> None: + """Require a resolver whenever a custom sink stores lineage. + + A missing resolver makes every continuation unresolved. + Current reconstruction refuses to guess across that boundary. + """ + if block.lineage_store is not None or installed_lineage_store() is not 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." + ) + return + raise ValueError( + "token_id_capture has a custom sink but no lineage_store. 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." + ) + @property def enabled(self) -> bool: return self.token_id_capture.enabled @@ -169,6 +201,18 @@ def build_sink(self) -> TokenSink | None: return None return self._build_endpoint(target, self.token_id_capture.sink_kwargs, TokenSink, "sink") + def build_lineage_store(self) -> LineageStore | None: + """Construct the configured request-time lineage store.""" + target = self.token_id_capture.lineage_store + if not self.token_id_capture.enabled or target is None: + return None + return self._build_endpoint( + target, + self.token_id_capture.lineage_store_kwargs, + LineageStore, + "lineage_store", + ) + @staticmethod def _build_endpoint(target: str, kwargs: dict[str, Any], protocol: type, kind: str): if ":" not in target: diff --git a/nemo_gym/token_id_capture/conformance.py b/nemo_gym/token_id_capture/conformance.py new file mode 100644 index 0000000000..959f3d60c5 --- /dev/null +++ b/nemo_gym/token_id_capture/conformance.py @@ -0,0 +1,320 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Verify an external token-capture transport against the ``protocols`` contracts. + +A framework transport (e.g. NeMo-RL over TransferQueue) replaces Gym's file store. +The contract points it must satisfy fail silently in production. +``run_conformance`` exercises each one directly. +Factories must return fresh client instances over the same shared backend. +Fresh instances are how cross-client visibility gets tested. +Each check uses its own rollout id, so failed checks cannot poison later ones. +Like ``protocols``, this module avoids FastAPI, Ray, Torch, and aiohttp imports. +""" + +from __future__ import annotations + +from typing import Awaitable, Callable + +from nemo_gym.token_id_capture.lineage import stamp_continuation +from nemo_gym.token_id_capture.protocols import LineageStore, TokenSink, TokenSource +from nemo_gym.token_id_capture.records import ( + ParentResolutionStatus, + TokenEntry, + cumulative_tokens, + stamp_lineage, +) + + +class ConformanceError(AssertionError): + """One named contract check failed.""" + + def __init__(self, check_name: str, detail: str) -> None: + self.check_name = check_name + self.detail = detail + super().__init__(f"{check_name}: {detail}") + + +def _require(condition: bool, check_name: str, detail: str) -> None: + if not condition: + raise ConformanceError(check_name, detail) + + +def _make_entry( + rollout_id: str, + model_call_id: str, + *, + prompt: list[int], + generation: list[int], + request_items: list[dict], + text: str, +) -> TokenEntry: + """Build a realistic committed entry, stamped exactly as ``capture_tokens`` would stamp it.""" + entry = TokenEntry( + rollout_id=rollout_id, + model_call_id=model_call_id, + model="conformance-model", + prompt_token_ids=list(prompt), + generation_token_ids=list(generation), + generation_log_probs=[-0.1] * len(generation), + output_items=[{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}], + token_item_index=0, + # Fixed so an idempotent retry can resend byte-identical payloads. + created_at=1700000000.0, + ) + stamp_continuation(entry, list(request_items)) + stamp_lineage(entry, None, parent_resolution=ParentResolutionStatus.ROOT) + return entry + + +def _identical_retry(entry: TokenEntry) -> TokenEntry: + """A retry resends the same serialized bytes; simulate one via a lossless round trip.""" + return TokenEntry.model_validate(entry.model_dump(mode="json")) + + +_REQUEST = [{"role": "user", "content": "What is the weather in Paris?"}] + + +async def run_conformance( + sink_factory: Callable[[], TokenSink], + source_factory: Callable[[], TokenSource], + lineage_factory: Callable[[], LineageStore] | None = None, + *, + rollout_id: str = "conformance-rollout", +) -> list[str]: + """Run the ordered contract checks and return the names that passed. + + Raise ``ConformanceError`` on the first failure. + Lineage checks are skipped without a ``lineage_factory``. + The ``begin_call`` check is skipped when the sink lacks the extension. + """ + passed: list[str] = [] + closables: list = [] + + def sink() -> TokenSink: + instance = sink_factory() + closables.append(instance) + return instance + + def source() -> TokenSource: + instance = source_factory() + closables.append(instance) + return instance + + def lineage() -> LineageStore: + assert lineage_factory is not None + instance = lineage_factory() + closables.append(instance) + return instance + + checks: list[tuple[str, Callable[[str], Awaitable[None]]]] = [ + ("put_then_freeze_visibility", lambda r: _check_put_then_freeze(sink(), source(), r)), + ("idempotent_reput", lambda r: _check_idempotent_reput(sink(), source(), r)), + ("conflicting_reput", lambda r: _check_conflicting_reput(sink(), source(), r)), + ("mark_incomplete_durability", lambda r: _check_mark_incomplete(sink(), source, r)), + ("freeze_idempotency", lambda r: _check_freeze_idempotency(sink(), source, r)), + ("post_freeze_write_safety", lambda r: _check_post_freeze_write(sink(), source(), r)), + ("conditional_retirement", lambda r: _check_conditional_retirement(sink(), source(), r)), + ] + if lineage_factory is not None: + checks.append( + ("lineage_visibility", lambda r: _check_lineage_visibility(sink(), lineage, r, fresh_client=False)) + ) + checks.append( + ( + "fresh_client_lineage_visibility", + lambda r: _check_lineage_visibility(sink(), lineage, r, fresh_client=True), + ) + ) + probe_sink = sink() + if getattr(probe_sink, "begin_call", None) is not None: + checks.append(("begin_call_custody", lambda r: _check_begin_call_custody(sink(), source(), r))) + + try: + for name, check in checks: + try: + await check(f"{rollout_id}-{name.replace('_', '-')}") + except ConformanceError: + raise + except Exception as error: # noqa: BLE001 - a raising backend is a conformance failure, not a crash. + raise ConformanceError(name, f"unexpected {type(error).__name__}: {error}") from error + passed.append(name) + finally: + for closable in closables: + await closable.close() + return passed + + +async def _check_put_then_freeze(sink: TokenSink, src: TokenSource, rollout_id: str) -> None: + name = "put_then_freeze_visibility" + entry = _make_entry(rollout_id, "call-1", prompt=[11, 12], generation=[13, 14], request_items=_REQUEST, text="a") + await sink.put(entry) + snapshot = await src.freeze(rollout_id) + _require(len(snapshot.entries) == 1, name, f"expected one entry, got {len(snapshot.entries)}") + _require(not snapshot.incomplete, name, "clean rollout froze incomplete") + _require(bool(snapshot.snapshot_id), name, "snapshot_id is empty") + frozen = snapshot.entries[0] + _require( + frozen.model_dump(mode="json") == entry.model_dump(mode="json"), + name, + "frozen entry does not round-trip the stored payload", + ) + + +async def _check_idempotent_reput(sink: TokenSink, src: TokenSource, rollout_id: str) -> None: + name = "idempotent_reput" + entry = _make_entry(rollout_id, "call-1", prompt=[11, 12], generation=[13, 14], request_items=_REQUEST, text="a") + await sink.put(entry) + await sink.put(_identical_retry(entry)) + snapshot = await src.freeze(rollout_id) + _require(len(snapshot.entries) == 1, name, f"byte-identical retry duplicated: {len(snapshot.entries)} entries") + _require(not snapshot.incomplete, name, "byte-identical retry marked the rollout incomplete") + + +async def _check_conflicting_reput(sink: TokenSink, src: TokenSource, rollout_id: str) -> None: + name = "conflicting_reput" + entry = _make_entry(rollout_id, "call-1", prompt=[11, 12], generation=[13, 14], request_items=_REQUEST, text="a") + conflict = _make_entry(rollout_id, "call-1", prompt=[11, 12], generation=[15], request_items=_REQUEST, text="b") + await sink.put(entry) + try: + await sink.put(conflict) + except Exception: + return # Fail-closed at the writer. + # Fail-closed at the reader: the rollout must not look trainable. + snapshot = await src.freeze(rollout_id) + _require( + snapshot.incomplete, + name, + "conflicting payload for one call id was accepted without raising or marking incomplete", + ) + + +async def _check_mark_incomplete(sink: TokenSink, source_factory: Callable[[], TokenSource], rollout_id: str) -> None: + name = "mark_incomplete_durability" + entry = _make_entry(rollout_id, "call-1", prompt=[11, 12], generation=[13, 14], request_items=_REQUEST, text="a") + await sink.put(entry) + before = await source_factory().freeze(rollout_id) + # After freeze: the marker must still land and must move the version. + await sink.mark_incomplete(rollout_id, "call-2") + after = await source_factory().freeze(rollout_id) + _require(after.incomplete, name, "a fresh source instance does not see the incomplete marker") + _require(after.version != before.version, name, "mark_incomplete did not change the observable version") + retired = await source_factory().drop(rollout_id, snapshot_id=before.snapshot_id, version=before.version) + _require(not retired, name, "a retirement staled by mark_incomplete succeeded") + + +async def _check_freeze_idempotency( + sink: TokenSink, source_factory: Callable[[], TokenSource], rollout_id: str +) -> None: + name = "freeze_idempotency" + entry = _make_entry(rollout_id, "call-1", prompt=[11, 12], generation=[13, 14], request_items=_REQUEST, text="a") + await sink.put(entry) + first = await source_factory().freeze(rollout_id) + second = await source_factory().freeze(rollout_id) + _require(first.snapshot_id == second.snapshot_id, name, "snapshot_id changed across idempotent freezes") + _require(first.version == second.version, name, "version changed across idempotent freezes") + _require( + {e.model_call_id: e.digest for e in first.entries} == {e.model_call_id: e.digest for e in second.entries}, + name, + "entry set changed across idempotent freezes", + ) + + +async def _check_post_freeze_write(sink: TokenSink, src: TokenSource, rollout_id: str) -> None: + name = "post_freeze_write_safety" + entry = _make_entry(rollout_id, "call-1", prompt=[11, 12], generation=[13, 14], request_items=_REQUEST, text="a") + await sink.put(entry) + snapshot = await src.freeze(rollout_id) + late = _make_entry( + rollout_id, "call-2", prompt=[11, 12, 13, 14], generation=[15], request_items=_REQUEST, text="b" + ) + try: + await sink.put(late) + except Exception: + return # A hard fence is acceptable. + # A relaxed fence must invalidate the consumed snapshot instead. + retired = await src.drop(rollout_id, snapshot_id=snapshot.snapshot_id, version=snapshot.version) + _require(not retired, name, "a write raced freeze, yet the stale snapshot was retired") + + +async def _check_conditional_retirement(sink: TokenSink, src: TokenSource, rollout_id: str) -> None: + name = "conditional_retirement" + entry = _make_entry(rollout_id, "call-1", prompt=[11, 12], generation=[13, 14], request_items=_REQUEST, text="a") + await sink.put(entry) + snapshot = await src.freeze(rollout_id) + stale = await src.drop(rollout_id, snapshot_id=snapshot.snapshot_id, version=snapshot.version + 1) + _require(not stale, name, "a stale-version retirement succeeded") + evidence = await src.freeze(rollout_id) + _require(len(evidence.entries) == 1, name, "a failed retirement discarded the evidence") + retired = await src.drop(rollout_id, snapshot_id=snapshot.snapshot_id, version=snapshot.version) + _require(retired, name, "retiring the exact consumed snapshot failed") + + +async def _check_lineage_visibility( + sink: TokenSink, + lineage_factory: Callable[[], LineageStore], + rollout_id: str, + *, + fresh_client: bool, +) -> None: + name = "fresh_client_lineage_visibility" if fresh_client else "lineage_visibility" + # The non-fresh variant resolves through a store that existed before the write. + store = lineage_factory() if not fresh_client else None + entry = _make_entry( + rollout_id, + "call-1", + prompt=[101, 102, 103], + generation=[201, 202], + request_items=_REQUEST, + text="It is sunny.", + ) + await sink.put(entry) + # The fresh variant constructs its resolver only after the entry is durable. + if store is None: + store = lineage_factory() + faithful = list(_REQUEST) + list(entry.output_items) + [{"role": "user", "content": "And tomorrow?"}] + resolution = await store.resolve(rollout_id, faithful) + _require( + resolution.status == ParentResolutionStatus.RESOLVED, + name, + f"a faithful continuation resolved as {resolution.status}: {resolution.reason}", + ) + assert resolution.match is not None + _require(resolution.match.model_call_id == "call-1", name, "resolved to the wrong parent call") + _require( + list(resolution.match.cumulative_token_ids) == cumulative_tokens(entry), + name, + "resolved parent carries the wrong cumulative tokens", + ) + rewritten = [{"role": "user", "content": "REWRITTEN"}] + list(entry.output_items) + rewritten_resolution = await store.resolve(rollout_id, rewritten) + _require( + rewritten_resolution.status == ParentResolutionStatus.UNRESOLVED, + name, + f"a rewritten context resolved as {rewritten_resolution.status}", + ) + root = await store.resolve(rollout_id, [{"role": "user", "content": "fresh question"}]) + _require( + root.status == ParentResolutionStatus.ROOT, + name, + f"a request without model history resolved as {root.status}", + ) + + +async def _check_begin_call_custody(sink: TokenSink, src: TokenSource, rollout_id: str) -> None: + name = "begin_call_custody" + await sink.begin_call(rollout_id, "call-lost") # type: ignore[attr-defined] + snapshot = await src.freeze(rollout_id) + _require(snapshot.incomplete, name, "a dangling pre-dispatch intent did not mask the rollout") diff --git a/nemo_gym/token_id_capture/consumer.py b/nemo_gym/token_id_capture/consumer.py index dfd9a1222c..977a0d7f44 100644 --- a/nemo_gym/token_id_capture/consumer.py +++ b/nemo_gym/token_id_capture/consumer.py @@ -19,7 +19,6 @@ Gym reads a frozen snapshot from the local token store. A trainer freezes the ``TokenSource`` provided by its transport. Both paths pass snapshot entries through the same build and projection. -Single-response delivery rejects ``per_request`` because it can return multiple trajectories. This module does not import rollout-record or model-server modules. The caller supplies the ``rollout_id``. @@ -97,14 +96,6 @@ def _assemble( builder: str, model: str, ) -> dict: - if builder == "per_request": - # Single-response delivery cannot represent multiple trajectories. - return _failed_build( - rollout_id, - builder, - "per_request returns multiple trajectories and is not supported by single-response delivery", - n_calls=len(entries), - ) # Mask a malformed rollout instead of failing the caller. # The contiguity check and projection can raise. # An uncaught exception could fail a full rollout or training batch. @@ -140,6 +131,8 @@ def _assemble( "delivered_fraction": notes.delivered_fraction, "generated_tokens_captured": notes.generated_tokens_captured, "generated_tokens_delivered": notes.generated_tokens_delivered, + "parent_link_failures": dict(notes.parent_link_failures), + "unresolved_parent_calls": len(notes.unresolved_parent_calls), # Calls without generated tokens have no training signal. # A nonzero count can indicate an output-budget or content-filter cutoff. "empty_generation_calls": len(notes.empty_generation_calls), @@ -152,8 +145,14 @@ def _assemble( "metrics": metrics, # A retry of the final call can leave two plausible generations. # Mask the rollout when the client-selected generation is unknown. - "mask_sample": bool(unresolved) or notes.roots != 1 or notes.chains != 1, + # An empty delivery must never be trainable, whatever produced it. + "mask_sample": bool(unresolved) + or bool(notes.unresolved_parent_calls) + or notes.roots != 1 + or notes.chains != 1 + or not any(item.get("generation_token_ids") for item in response.get("output", [])), "unresolved_retries": list(unresolved), + "unresolved_parent_calls": list(notes.unresolved_parent_calls), } diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py new file mode 100644 index 0000000000..32285070cf --- /dev/null +++ b/nemo_gym/token_id_capture/lineage.py @@ -0,0 +1,713 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Resolve the recorded call that a request continues. + +A rollout can contain several model calls. +Training consumes their exact tokens as one contiguous sequence. +Request-time lineage identifies the earlier call that each request continues. + +``assistant_fingerprint`` is the lookup key. +It hashes model-authored turns and ignores user and tool content added between calls. +``conversation_digest`` verifies the unchanged request context. +A digest mismatch rejects the claimed lineage before any parent tokens are reused. + +The shared ``LineageStore`` resolves entries already committed by ``TokenSink``. +``FileLineageStore`` tails the token JSONL through the token store's lock. +Each child receives its parent's cumulative tokens. +Downstream inference consumes those tokens to supply the exact prompt prefix. + +Every new record distinguishes a root, a resolved parent, and an unresolved boundary. +Only records that predate this metadata use token-prefix fallback. + +The guaranteed invariant is token-chain exactness, not conversation fidelity. +A delivered chain contains exactly the tokens the policy emitted over the recorded context. +The hashes deliberately ignore reasoning and some inserted items. +Those fields may differ from the harness rendering without breaking token-chain exactness. +Ambiguous matches remain unresolved rather than risking tokens from the wrong call. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import orjson + +from nemo_gym.token_id_capture.protocols import LineageMatch, LineageResolution +from nemo_gym.token_id_capture.records import ParentResolutionStatus, TokenEntry, cumulative_tokens + + +# Increment when fingerprint canonicalization or hash layout changes. +# Resolvers ignore entries stamped with a different version. +FINGERPRINT_VERSION = 1 + +_FINGERPRINT_DOMAIN = b"nemo-gym-lineage" +_CONTEXT_DOMAIN = b"nemo-gym-lineage-context" + + +def _update_field(hasher: Any, tag: bytes, value: str) -> None: + """Hash one tagged, length-delimited UTF-8 field.""" + encoded = value.encode("utf-8") + hasher.update(tag) + hasher.update(len(encoded).to_bytes(8, "big")) + hasher.update(encoded) + + +def _canonical_json(value: Any) -> str: + """Serialize JSON-compatible prompt content without losing structure.""" + try: + return orjson.dumps(value, option=orjson.OPT_SORT_KEYS).decode("utf-8") + except (TypeError, orjson.JSONEncodeError) as error: + raise ValueError(f"unsupported prompt content: {type(value).__name__}") from error + + +def canonicalize_tool_arguments(value: Any) -> str: + """Normalize a tool call's arguments for comparison only. + + Harnesses can reserialize tool-call arguments between turns. + Comparison uses sorted-key JSON with normalized separators. + The record retains the model's original string. + """ + if value is None: + return "" + if isinstance(value, str): + try: + parsed = json.loads(value) + except (TypeError, ValueError): + return value.strip() + else: + parsed = value + return _canonical_json(parsed) + + +def _content_of(content: Any) -> list[tuple[str, str]]: + """Return typed content parts without discarding prompt-shaping blocks. + + Tool calls are normalized separately by ``_tools_of``. + Tool results are normalized separately by ``_tool_results_of``. + """ + if content is None: + return [] + if isinstance(content, str): + return [("text", content)] if content else [] + if not isinstance(content, list): + raise ValueError(f"unsupported message content: {type(content).__name__}") + parts: list[tuple[str, str]] = [] + for block in content: + if isinstance(block, str): + if block: + parts.append(("text", block)) + continue + if not isinstance(block, dict): + raise ValueError(f"unsupported content block: {type(block).__name__}") + block_type = str(block.get("type") or "") + if block_type in {"tool_use", "tool_result"}: + continue + if isinstance(block.get("text"), str) and block_type in { + "", + "text", + "input_text", + "output_text", + }: + if block["text"]: + parts.append(("text", block["text"])) + continue + if not block_type: + raise ValueError("content block has no supported type") + parts.append((block_type, _canonical_json(block))) + return parts + + +def _tools_of(message: dict) -> list[tuple[str, str, str]]: + """Return tool calls as ``(id, name, canonical arguments)`` tuples. + + Chat stores calls in the message's ``tool_calls`` field. + Anthropic stores calls in ``tool_use`` content blocks. + Responses stores each call as a standalone ``function_call`` item. + """ + tools: list[tuple[str, str, str]] = [] + # A Responses item is the tool call. + if message.get("type") == "function_call": + tools.append( + ( + str(message.get("call_id") or message.get("id") or ""), + str(message.get("name", "")), + canonicalize_tool_arguments(message.get("arguments")), + ) + ) + content = message.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_use": + tools.append( + ( + str(block.get("id") or ""), + str(block.get("name", "")), + canonicalize_tool_arguments(block.get("input")), + ) + ) + for call in message.get("tool_calls") or []: + function = (call or {}).get("function") or {} + tools.append( + ( + str((call or {}).get("id") or ""), + str(function.get("name", "")), + canonicalize_tool_arguments(function.get("arguments")), + ) + ) + return tools + + +def _tool_results_of(message: dict) -> list[tuple[str, str]]: + """Return tool result identities and payloads across dialects. + + Responses stores results in standalone ``function_call_output`` items. + Anthropic stores results in ``tool_result`` content blocks. + Chat stores results as plain message content. + """ + parts: list[tuple[str, str]] = [] + if message.get("type") == "function_call_output": + output = message.get("output") + parts.append( + ( + str(message.get("call_id") or message.get("id") or ""), + output if isinstance(output, str) else _canonical_json(output), + ) + ) + elif message.get("role") == "tool": + content = message.get("content") + parts.append( + ( + str(message.get("tool_call_id") or ""), + content if isinstance(content, str) else _canonical_json(content), + ) + ) + content = message.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_result": + inner = block.get("content") + payload = inner if isinstance(inner, str) else _canonical_json(inner) + parts.append((str(block.get("tool_use_id") or block.get("id") or ""), payload)) + return parts + + +def _is_assistant_authored(message: dict) -> bool: + """Return whether the model produced this item. + + Chat and Anthropic use the ``assistant`` role. + Responses tool calls are roleless ``function_call`` items. + """ + if message.get("role") == "assistant": + return True + # Reasoning is deliberately excluded. + # A harness need not echo standalone reasoning items. + # Including reasoning would make fingerprints depend on the dialect and echo behavior. + # Reasoning-only collisions resolve as ambiguous and fall back. + return message.get("type") == "function_call" + + +def conversation_digest(messages: list[dict]) -> str: + """Hash every turn of a conversation, model-authored or not. + + ``assistant_fingerprint`` ignores user and tool content. + This digest covers that omitted context. + A mismatch rejects the parent before its tokens are reused. + """ + hasher = hashlib.sha256(_CONTEXT_DOMAIN) + for message in messages or []: + if not isinstance(message, dict): + raise ValueError(f"request item is not an object: {type(message).__name__}") + _update_field(hasher, b"\x00", str(message.get("role") or message.get("type") or "")) + for content_type, payload in _content_of(message.get("content")): + _update_field(hasher, b"\x01", content_type) + _update_field(hasher, b"\x02", payload) + for call_id, name, arguments in _tools_of(message): + _update_field(hasher, b"\x03", call_id) + _update_field(hasher, b"\x04", name) + _update_field(hasher, b"\x05", arguments) + # Include tool results. + # Summarizing, redacting, or truncating a result changes the request context. + for call_id, output in _tool_results_of(message): + _update_field(hasher, b"\x06", call_id) + _update_field(hasher, b"\x07", output) + return hasher.hexdigest() + + +def assistant_fingerprint(messages: list[dict]) -> str: + """Fingerprint the model-authored turns of a request, in order. + + The fingerprint identifies the call that produced the last model-authored turn. + User and tool content is excluded from the lookup key. + Dialect-specific tool-call shapes normalize to the same hash input. + """ + hasher = hashlib.sha256(_FINGERPRINT_DOMAIN) + count = 0 + for message in messages or []: + if not isinstance(message, dict): + raise ValueError(f"request item is not an object: {type(message).__name__}") + if not _is_assistant_authored(message): + continue + count += 1 + for content_type, payload in _content_of(message.get("content")): + _update_field(hasher, b"\x00", content_type) + _update_field(hasher, b"\x01", payload) + for call_id, name, arguments in _tools_of(message): + _update_field(hasher, b"\x02", call_id) + _update_field(hasher, b"\x03", name) + _update_field(hasher, b"\x04", arguments) + if count == 0: + return "" + return hasher.hexdigest() + + +@dataclass +class LineageNode: + call_id: str + # ``None`` means the index is metadata-only. + # A resolved match loads tokens from ``entry_offset``. + cum_tokens: list[int] | None + cum_len: int + digest: str + entry_offset: int = -1 + # These fields describe the request context sent for this call. + # They exclude the model's response. + # The item count is stable while the harness stays in one dialect. + # A mid-rollout dialect switch can misalign it; verification then fails closed. + context_len: int = 0 + context_digest: str = "" + + +def stamp_continuation(entry: TokenEntry, request_items: list[dict]) -> TokenEntry: + """Add compact lookup metadata before the token entry is committed.""" + entry.continuation_fingerprint = assistant_fingerprint(list(request_items) + list(entry.output_items)) + entry.continuation_context_len = len(request_items) + entry.continuation_context_digest = conversation_digest(request_items) + entry.fingerprint_version = FINGERPRINT_VERSION + return entry + + +@dataclass +class RolloutLineage: + """Keep an append-only per-rollout call index.""" + + by_fingerprint: dict[str, list[str]] = field(default_factory=dict) + by_call_id: dict[str, LineageNode] = field(default_factory=dict) + # Cache the cumulative token count for memory bounds. + total_tokens: int = 0 + + def resolve_node(self, messages: list[dict]) -> tuple[ParentResolutionStatus, "LineageNode | None", str]: + """Return the parent decision without touching token arrays. + + Matching needs only fingerprints, digests, and lengths. + The caller materializes tokens for the single winner. + """ + fingerprint = assistant_fingerprint(messages) + if not fingerprint: + return ParentResolutionStatus.ROOT, None, "" + # dict.fromkeys: a call id indexed twice (e.g. by racing refreshes) is one candidate. + call_ids = list(dict.fromkeys(self.by_fingerprint.get(fingerprint) or [])) + candidates = [ + node + for call_id in call_ids + if (node := self.by_call_id.get(call_id)) is not None and self._continues(node, messages) + ] + if len(candidates) > 1: + # Calls with identical cumulative tokens are interchangeable. + # Keep different token sequences unresolved. + digests = {(node.digest, node.cum_len) for node in candidates} + if len(digests) == 1 and candidates[0].digest: + candidates = [min(candidates, key=lambda node: node.call_id)] + if len(candidates) != 1: + return ParentResolutionStatus.UNRESOLVED, None, "no_match" if not candidates else "ambiguous" + return ParentResolutionStatus.RESOLVED, candidates[0], "" + + def resolve(self, messages: list[dict]) -> LineageResolution: + """Return the immutable parent decision for this request. + + A request without model-authored history is a root. + A request with unverified history is unresolved. + Never guess among calls with identical output. + """ + status, node, reason = self.resolve_node(messages) + if status != ParentResolutionStatus.RESOLVED: + return LineageResolution(status, reason=reason) + if node.cum_tokens is None: + raise ValueError("metadata-only lineage node requires caller-side materialization") + return LineageResolution( + ParentResolutionStatus.RESOLVED, + match=LineageMatch( + model_call_id=node.call_id, + cumulative_token_ids=tuple(node.cum_tokens), + digest=node.digest, + ), + ) + + @staticmethod + def _continues(node: LineageNode, messages: list[dict]) -> bool: + """Return whether this request extends the node's recorded context. + + The leading ``context_len`` items must match the recorded request. + A rewritten or summarized context fails verification. + Verification excludes the model response because dialects can echo it as different item counts. + """ + if not node.context_digest: + # Fail closed when no context digest is available. + return False + if len(messages) < node.context_len: + return False + return conversation_digest(messages[: node.context_len]) == node.context_digest + + def add_entry(self, entry: TokenEntry, *, store_tokens: bool = True, entry_offset: int = -1) -> None: + """Index lookup metadata carried by one committed token entry. + + ``store_tokens=False`` keeps token arrays in the durable log. + """ + if not entry.continuation_fingerprint: + return + if entry.fingerprint_version is not None and entry.fingerprint_version != FINGERPRINT_VERSION: + # A different algorithm produced this fingerprint; matching it would be luck. + return + if getattr(entry, "prompt_is_delta", False) and store_tokens: + # A memory-only index cannot reconstruct a delta chain. + raise ValueError("delta records require a durable-log-backed lineage store") + node = LineageNode( + call_id=entry.model_call_id, + cum_tokens=cumulative_tokens(entry) if store_tokens else None, + cum_len=entry.cum_len if entry.cum_len is not None else len(cumulative_tokens(entry)), + digest=entry.digest or "", + entry_offset=entry_offset, + context_len=entry.continuation_context_len, + context_digest=entry.continuation_context_digest, + ) + previous = self.by_call_id.get(entry.model_call_id) + if previous is not None: + if previous != node: + raise ValueError(f"conflicting lineage record for model call {entry.model_call_id}") + return + self.total_tokens += node.cum_len + self.by_call_id[entry.model_call_id] = node + self.by_fingerprint.setdefault(entry.continuation_fingerprint, []).append(entry.model_call_id) + + def record( + self, + call_id: str, + messages: list[dict], + cum_tokens: list[int], + digest: str, + context_len: int | None = None, + ) -> None: + """Build an in-memory entry for direct index tests.""" + request_len = context_len if context_len is not None else max(len(messages) - 1, 0) + entry = TokenEntry( + rollout_id="_in_memory", + model_call_id=call_id, + prompt_token_ids=[], + generation_token_ids=list(cum_tokens), + generation_log_probs=[0.0] * len(cum_tokens), + output_items=list(messages[request_len:]), + cum_len=len(cum_tokens), + digest=digest, + continuation_fingerprint=assistant_fingerprint(messages), + continuation_context_len=request_len, + continuation_context_digest=conversation_digest(messages[:request_len]), + ) + self.add_entry(entry) + + +class LineageIndex: + """Bound worker-local lineage by rollout and cumulative token counts. + + This index backs the single-worker fallback. + Shared stores provide cross-worker visibility. + Eviction removes the oldest rollout. + An evicted parent leaves later continuations unresolved and the builder masks them. + The only live rollout is never evicted. + """ + + def __init__(self, max_rollouts: int = 512, max_tokens: int = 8_000_000) -> None: + self._max_rollouts = max_rollouts + self._max_tokens = max_tokens + self._rollouts: dict[str, RolloutLineage] = {} + + def for_rollout(self, rollout_id: str) -> RolloutLineage: + lineage = self._rollouts.get(rollout_id) + if lineage is None: + lineage = RolloutLineage() + self._rollouts[rollout_id] = lineage + self._evict() + return lineage + + def _evict(self) -> None: + # Check after every access because existing rollouts can grow. + while self._rollouts and (len(self._rollouts) > self._max_rollouts or self.total_tokens > self._max_tokens): + oldest = next(iter(self._rollouts)) + # Never evict the only rollout. + if len(self._rollouts) == 1: + return + self._rollouts.pop(oldest) + + @property + def total_tokens(self) -> int: + return sum(lineage.total_tokens for lineage in self._rollouts.values()) + + def drop(self, rollout_id: str) -> None: + """Release a rollout's lineage early. + + Gym's model server has no rollout-completion signal. + An in-process framework can call this when it retires the records. + """ + self._rollouts.pop(rollout_id, None) + + def clear(self) -> None: + self._rollouts.clear() + + def __len__(self) -> int: + return len(self._rollouts) + + +class InMemoryLineageStore: + """Reference resolver for in-process framework backends and tests. + + Production wiring uses ``FileLineageStore`` when a token store exists. + This class supports in-process framework adapters and tests. + Its index is memory-only. + Eviction or restart leaves affected continuations unresolved. + That failure mode is safe but can mask otherwise usable rollouts. + Production adapters should back the incremental resolver with durable records. + """ + + def __init__(self, max_rollouts: int = 512, max_tokens: int = 8_000_000) -> None: + self.index = LineageIndex(max_rollouts=max_rollouts, max_tokens=max_tokens) + + async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: + return self.index.for_rollout(rollout_id).resolve(request_items) + + async def put(self, entry: TokenEntry) -> None: + """Publish one committed entry to the worker-local index.""" + self.index.for_rollout(entry.rollout_id).add_entry(entry) + + def is_process_shared(self) -> bool: + return False + + async def close(self) -> None: + self.index.clear() + + +class IncrementalLineageStore: + """Base class for lineage resolvers over any committed-entry backend. + + An external backend implements two hooks. + It inherits Gym's matcher, bounded index, locking, and token materialization. + Hash-for-hash agreement is the wire contract. + The backend remains the source of truth when cache rows are evicted. + A resolved match loads only the winning call's token chain. + + Required hooks: + ``_fetch_new_entries(rollout_id, cursor)`` -> ``(items, new_cursor)`` where + ``items`` is ``[(TokenEntry, ref), ...]`` in commit order since ``cursor`` + (``None`` means from the beginning) and ``ref`` is any handle that + ``_load_entry`` can use later (byte offset, KV key, ...). Raise + ``CursorReset`` when the cursor no longer describes the backend (file + rotated, namespace recreated); the base refetches from the beginning. + ``_load_entry(rollout_id, ref)`` -> ``TokenEntry`` for one committed record. + + Optional hooks: + ``_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 + be shared, and the multi-worker startup check trusts this answer. + """ + + class CursorReset(Exception): + """The stored cursor no longer describes the backend; refetch from scratch.""" + + def __init__(self, *, max_cached_rollouts: int = 65536) -> None: + import threading + + if max_cached_rollouts < 1: + raise ValueError("max_cached_rollouts 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 + self._cache_guard = threading.Lock() + # Fixed lock striping bounds synchronization metadata. + # Hash collisions only serialize unrelated rollouts. + self._rollout_locks = tuple(threading.Lock() for _ in range(256)) + + # -- hooks ---------------------------------------------------------------- + def _fetch_new_entries(self, rollout_id: str, cursor: Any) -> tuple[list[tuple[TokenEntry, Any]], Any]: + raise NotImplementedError + + def _load_entry(self, rollout_id: str, ref: Any) -> TokenEntry: + raise NotImplementedError + + def _read_locked(self, rollout_id: str): + from contextlib import nullcontext + + return nullcontext() + + # -- shared machinery ----------------------------------------------------- + def _rollout_lock(self, rollout_id: str): + return self._rollout_locks[hash(rollout_id) % len(self._rollout_locks)] + + def _cache_put(self, rollout_id: str, value: tuple[Any, dict[str, Any], RolloutLineage]) -> None: + """Insert or touch a cache row with LRU semantics. + + Reinsert a touched row so dictionary order tracks recency. + Eviction only requires a later backend refetch. + """ + with self._cache_guard: + self._cache.pop(rollout_id, None) + self._cache[rollout_id] = value + while len(self._cache) > self._max_cached_rollouts: + oldest = next(iter(self._cache)) + if oldest == rollout_id: + break + self._cache.pop(oldest) + + def _refresh(self, rollout_id: str) -> tuple[dict[str, Any], RolloutLineage]: + with self._cache_guard: + cached = self._cache.get(rollout_id) + cursor, refs, lineage = cached if cached is not None else (None, {}, RolloutLineage()) + try: + items, cursor = self._fetch_new_entries(rollout_id, cursor) + except IncrementalLineageStore.CursorReset: + refs, lineage = {}, RolloutLineage() + items, cursor = self._fetch_new_entries(rollout_id, None) + for entry, ref in items: + refs[entry.model_call_id] = ref + # Metadata-only: tokens stay in the backend behind ``ref``. + lineage.add_entry(entry, store_tokens=False, entry_offset=ref if isinstance(ref, int) else -1) + self._cache_put(rollout_id, (cursor, refs, lineage)) + return refs, lineage + + def _materialize( + self, rollout_id: str, node: LineageNode, refs: dict[str, Any], lineage: RolloutLineage + ) -> list[int]: + """Load one RESOLVED parent's cumulative tokens from the backend. + + 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: + 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 + if node.digest and compute_digest(tokens) != node.digest: + raise ValueError(f"materialized tokens for {node.call_id} fail their digest") + return tokens + + async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: + return await asyncio.to_thread(self._resolve, rollout_id, request_items) + + def _resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: + with self._rollout_lock(rollout_id), self._read_locked(rollout_id): + refs, lineage = self._refresh(rollout_id) + status, node, reason = lineage.resolve_node(request_items) + if status != ParentResolutionStatus.RESOLVED: + return LineageResolution(status, reason=reason) + tokens = ( + node.cum_tokens if node.cum_tokens is not None else self._materialize(rollout_id, node, refs, lineage) + ) + return LineageResolution( + ParentResolutionStatus.RESOLVED, + match=LineageMatch( + model_call_id=node.call_id, + cumulative_token_ids=tuple(tokens), + digest=node.digest, + ), + ) + + def is_process_shared(self) -> bool: + return True + + async def close(self) -> None: + with self._cache_guard: + self._cache.clear() + + +class FileLineageStore(IncrementalLineageStore): + """Resolve lineage from the token JSONL committed by ``TokenCaptureStore``. + + The reference ``IncrementalLineageStore`` backend: cursor = (inode, offset), + ref = byte offset, reads under the store's shared flock so a committed + ``put`` is immediately visible. + """ + + def __init__(self, root: str | Path, *, max_cached_rollouts: int = 65536) -> None: + from nemo_gym.token_id_capture.store import TokenCaptureStore + + super().__init__(max_cached_rollouts=max_cached_rollouts) + self._store = TokenCaptureStore(root) + + def _read_locked(self, rollout_id: str): + return self._store._locked(rollout_id, shared=True) + + def _fetch_new_entries(self, rollout_id: str, cursor: Any) -> tuple[list[tuple[TokenEntry, Any]], Any]: + path = self._store.path_for(rollout_id) + if not path.exists(): + if cursor is not None: + raise IncrementalLineageStore.CursorReset + return [], None + file_stat = path.stat() + inode, offset = cursor if cursor is not None else (file_stat.st_ino, 0) + if inode != file_stat.st_ino or offset < 0 or offset > file_stat.st_size: + raise IncrementalLineageStore.CursorReset + items: list[tuple[TokenEntry, Any]] = [] + if offset < file_stat.st_size: + with path.open("rb") as handle: + handle.seek(offset) + while True: + line_offset = handle.tell() + line = handle.readline() + if not line: + break + payload = line.strip() + if not payload: + continue + items.append((TokenEntry.model_validate(orjson.loads(payload)), line_offset)) + offset = handle.tell() + return items, (inode, offset) + + 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())) diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index 20c7637905..a5b2d36610 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -23,6 +23,22 @@ Consumers read through ``TokenSource.freeze``. They identify the frozen state with ``snapshot_id``. This module avoids FastAPI, Ray, Torch, and aiohttp imports. + +The load-bearing guarantee is the happens-before edge through ``TokenSink.put``. +``put`` is awaited on the serving path. +The harness therefore cannot send a continuation before the previous call's record is durable. +The record must also be visible to any worker's lineage resolver when ``put`` returns. +A transport that acknowledges ``put`` before cross-client visibility breaks this silently. +The break appears only as a load-dependent trickle of unresolved, masked samples. + +A sink may additionally implement ``begin_call(rollout_id, model_call_id)``. +It is an optional extension and deliberately not part of the ``TokenSink`` protocol. +``begin_call`` durably records a pre-dispatch intent. +An intent with no matching entry at freeze must mask the rollout. +That closes the window where the final call's entry is lost without a trace. +``begin_call`` runs before generation, so the caller may fail the model call at zero compute cost. + +``nemo_gym.token_id_capture.conformance`` checks an external implementation against these contracts. """ from __future__ import annotations @@ -30,7 +46,7 @@ from dataclasses import dataclass from typing import Protocol, runtime_checkable -from nemo_gym.token_id_capture.records import TokenEntry +from nemo_gym.token_id_capture.records import ParentResolutionStatus, TokenEntry @dataclass(frozen=True) @@ -44,16 +60,86 @@ class TokenCaptureSnapshot: version: int +@dataclass(frozen=True) +class LineageMatch: + """Describe a uniquely verified parent from a shared lineage store.""" + + model_call_id: str + cumulative_token_ids: tuple[int, ...] + digest: str + + +@dataclass(frozen=True) +class LineageResolution: + """Return one immutable request-time parent decision.""" + + status: ParentResolutionStatus + match: LineageMatch | None = None + reason: str = "" + + def __post_init__(self) -> None: + if self.status == ParentResolutionStatus.RESOLVED and self.match is None: + raise ValueError("resolved lineage requires a match") + if self.status != ParentResolutionStatus.RESOLVED and self.match is not None: + raise ValueError(f"{self.status.value} lineage cannot carry a match") + + +@runtime_checkable +class LineageStore(Protocol): + """Resolve request-time lineage from entries committed by a token sink. + + This is a read-only view over sink-committed records. + After ``TokenSink.put`` returns, a later ``resolve`` on any worker must see the entry. + Visibility is required per rollout key only, so the store may shard by rollout. + Implementations must never guess among candidates. + ``resolve`` may fail toward UNRESOLVED on any doubt. + The offline builder independently re-verifies every claimed link by digest. + That backstop is what makes the relaxations in this module safe. + External implementations should embed Gym's ``RolloutLineage`` matcher rather than reimplement the hashing. + ``nemo_gym.token_id_capture.lineage`` is importable without the server stack. + Callers of ``commit_entry`` outside ``capture_tokens`` must run ``stamp_continuation`` themselves. + Unstamped entries are invisible to resolution. + """ + + async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: + """Return whether the request is a root, resolved, or unresolved. + + ``request_items`` are the unmodified harness items. + The implementation must verify the recorded request context. + A conflicting set of committed payloads for one call id must count as zero candidates. + UNRESOLVED is always a safe answer; a wrong RESOLVED is caught later by digest verification. + """ + ... + + def is_process_shared(self) -> bool: + """Return whether separate model-server workers share committed entries.""" + ... + + async def close(self) -> None: + """Release resources. Idempotent. + + The store is read-only, so there is never pending work to flush. + """ + ... + + @runtime_checkable class TokenSink(Protocol): """Receive captured records through Gym's file store or a framework transport.""" async def put(self, entry: TokenEntry) -> None: - """Durably store one record. + """Durably store one record before returning. + The entry carries its continuation lookup metadata. + Durability and resolver visibility are the return condition, not an eventual goal. + Any worker's paired lineage resolver must see the entry once this method returns. Repeating the same call id with the same payload is a no-op. + "Same payload" means the identical serialized entry, byte-for-byte, timestamps included. + A retry must resend the same bytes; rebuilding the entry produces a conflict, not a retry. Reusing a call id with a different payload must fail. - Writing after the rollout is frozen must fail. + A transport without compare-and-swap may delegate that conflict to the reader. + A resolver must then treat conflicting committed payloads for one call id as zero candidates. + Writing after freeze must fail or bump the frozen version; see ``TokenSource.freeze``. This method may raise. The caller marks the rollout incomplete. @@ -68,11 +154,20 @@ async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> Non A consumer must mask the sample instead of training on a chain with a hole. The model call itself still succeeds. This marker is therefore the durable signal that capture failed. + It must succeed after freeze. + It must change the observable version; that is what invalidates a stale retirement. + Make it more available than ``put``, for example through a local spill. + ``put`` and ``mark_incomplete`` failing together is the silent-loss case. """ ... async def close(self) -> None: - """Flush pending work and release resources idempotently.""" + """Release resources. Idempotent. + + There is never buffered unwritten data here. + ``put`` guaranteed durability before it returned. + A close that must flush records means ``put`` broke its contract. + """ ... @@ -84,8 +179,15 @@ async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: """Freeze a rollout and return one atomic snapshot. Freezing is idempotent. - No successful writes may occur after it returns. Entry order carries no meaning. + Entries are unique per ``model_call_id``. + An at-least-once transport must dedupe identical copies before snapshotting. + The fence is relaxed: "no successful write after freeze" is not required cluster-wide. + A strict fence is unimplementable without compare-and-swap. + A write racing freeze may therefore succeed durably. + It must then bump the version. + A conditional ``drop`` of the consumed snapshot then fails, and the evidence is retained. + Attempt-scoped rollout ids are the sanctioned strategy for retirement without compare-and-swap. """ ... @@ -93,8 +195,7 @@ async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool """Conditionally retire the exact frozen snapshot that was consumed. Return ``False`` if state changed after the snapshot. - Implementations that cannot delete return ``True``. - Their owner remains responsible for retention. + Transports without delete return ``True`` and own retention. """ ... @@ -108,10 +209,17 @@ async def close(self) -> None: # Request-scoped sinks take precedence. _INSTALLED_SINK: TokenSink | None = None _INSTALLED_SOURCE: TokenSource | None = None +_INSTALLED_LINEAGE_STORE: LineageStore | None = None def install_token_sink(sink: TokenSink | None) -> None: """Set (or clear, with ``None``) the process-wide default sink.""" + if sink is not None: + # A sink missing mark_incomplete can hide incomplete rollouts. + # Validate programmatically installed sinks too. + missing = [name for name in ("put", "mark_incomplete", "close") if not callable(getattr(sink, name, None))] + if missing: + raise TypeError(f"installed token sink lacks required methods: {', '.join(missing)}") global _INSTALLED_SINK _INSTALLED_SINK = sink @@ -131,3 +239,13 @@ def install_token_source(source: TokenSource | None) -> None: def installed_token_source() -> TokenSource | None: return _INSTALLED_SOURCE + + +def install_lineage_store(store: LineageStore | None) -> None: + """Set (or clear) the process-wide request-time lineage store.""" + global _INSTALLED_LINEAGE_STORE + _INSTALLED_LINEAGE_STORE = store + + +def installed_lineage_store() -> LineageStore | None: + return _INSTALLED_LINEAGE_STORE diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py index 7f77b3e725..594a59658d 100644 --- a/nemo_gym/token_id_capture/records.py +++ b/nemo_gym/token_id_capture/records.py @@ -25,6 +25,9 @@ from __future__ import annotations +import hashlib +import struct +from enum import StrEnum from typing import Any from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -41,7 +44,56 @@ # ``extra="allow"`` otherwise hides unknown fields. # # 1 rollout and call identity, the token arrays, the output items and their carrier index -TOKEN_ENTRY_RECORD_SCHEMA_VERSION = 1 +# 2 parent_call_id, cum_len and digest, added when calls began being linked to their parent +# 3 parent resolution and compact continuation lookup metadata +TOKEN_ENTRY_RECORD_SCHEMA_VERSION = 3 + +# Records below schema 3 never left development. +# Current records always carry a request-time parent decision. +TOKEN_ENTRY_MIN_SCHEMA_VERSION = 3 + +# Increment this version when the digest encoding changes. +# A stale digest must fail verification. +DIGEST_VERSION = 1 +_DIGEST_DOMAIN = b"nemo-gym-tokens" +_EMPTY_DIGEST = hashlib.sha256(_DIGEST_DOMAIN).hexdigest() + + +class ParentResolutionStatus(StrEnum): + """Describe whether a model call has a proven captured predecessor.""" + + ROOT = "root" + RESOLVED = "resolved" + UNRESOLVED = "unresolved" + + +def encode_token_ids(token_ids: list[int]) -> bytes: + """Stable, length-delimited, big-endian encoding of a token sequence. + + Independent implementations can hash the same bytes. + One vectorized pack replaces the per-token loop. + The byte layout remains identical to the original encoding. + """ + header = struct.pack(">BQ", DIGEST_VERSION, len(token_ids)) + if not token_ids: + return header + try: + return header + struct.pack(f">{len(token_ids)}Q", *token_ids) + except struct.error: + negative = next(token_id for token_id in token_ids if token_id < 0) + raise ValueError(f"token ids must be non-negative, got {negative}") from None + + +def compute_digest(token_ids: list[int]) -> str: + """Digest of an exact token sequence. + + The builder verifies a claimed parent by hashing the corresponding prompt prefix. + A mismatch quarantines the call. + This prevents stale or interleaved records from merging silently. + """ + if not token_ids: + return _EMPTY_DIGEST + return hashlib.sha256(_DIGEST_DOMAIN + encode_token_ids(token_ids)).hexdigest() class TokenEntry(BaseModel): @@ -74,6 +126,32 @@ class TokenEntry(BaseModel): # This non-semantic timestamp helps diagnose retries and sibling branches. created_at: float = 0.0 + # These fields were added with parent resolution in schema version 2. + # A missing parent call ID triggers strict token-prefix matching. + # + # A parent link identifies the exact call retained by the harness. + # This disambiguates retries with the same prompt. + # The builder verifies the link instead of trusting it. + parent_call_id: str | None = None + # This is the length of this call's prompt and generation. + # A child must start with a prefix of this length. + cum_len: int | None = None + # This is ``compute_digest(prompt_token_ids + generation_token_ids)``. + digest: str | None = None + # New records distinguish a valid root from a missing parent. + # ``None`` is reserved for records written before schema version 3. + parent_resolution: ParentResolutionStatus | None = None + # These fields make a committed entry visible to request-time resolution. + # The fingerprint identifies the model-authored output. + continuation_fingerprint: str = "" + # 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. + parent_resolution_reason: str = "" + # Identify the continuation fingerprint algorithm. + fingerprint_version: int | None = None + @model_validator(mode="after") def _refuse_a_newer_record(self) -> "TokenEntry": """Accept older records and reject newer records. @@ -88,6 +166,11 @@ def _refuse_a_newer_record(self) -> "TokenEntry": f"up to {TOKEN_ENTRY_RECORD_SCHEMA_VERSION}. Upgrade the reader, or point it at " "records written by a writer it matches." ) + 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." + ) if len(self.generation_token_ids) != len(self.generation_log_probs): raise ValueError( "generation_token_ids and generation_log_probs must have the same length " @@ -97,9 +180,43 @@ def _refuse_a_newer_record(self) -> "TokenEntry": raise ValueError( f"token_item_index {self.token_item_index} is outside output_items of length {len(self.output_items)}" ) + if self.parent_resolution == ParentResolutionStatus.RESOLVED and not self.parent_call_id: + raise ValueError("resolved parent metadata requires parent_call_id") + 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") return self +def cumulative_tokens(entry: TokenEntry) -> list[int]: + """The full sequence a child of this call must start with.""" + return list(entry.prompt_token_ids) + list(entry.generation_token_ids) + + +def stamp_lineage( + entry: TokenEntry, + parent_call_id: str | None, + *, + parent_resolution: ParentResolutionStatus | None = None, +) -> TokenEntry: + """Fill token lineage and the request-time parent decision. + + ``cum_len`` and ``digest`` always describe this call. + ``parent_resolution=None`` preserves records built by compatibility callers. + """ + cumulative = cumulative_tokens(entry) + entry.cum_len = len(cumulative) + entry.digest = compute_digest(cumulative) + entry.parent_call_id = parent_call_id + entry.parent_resolution = parent_resolution + if parent_resolution == ParentResolutionStatus.RESOLVED and parent_call_id is None: + raise ValueError("resolved parent metadata requires parent_call_id") + if parent_resolution in {ParentResolutionStatus.ROOT, ParentResolutionStatus.UNRESOLVED}: + if parent_call_id is not None: + raise ValueError(f"{parent_resolution.value} parent metadata cannot carry parent_call_id") + return entry + + def response_to_output_items(payload: dict) -> list[dict]: """Normalize a served response to a list of content-bearing Responses output items. diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index 69ddde7c7d..0c921ddf09 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -27,16 +27,20 @@ from __future__ import annotations import logging +import threading import time from contextvars import ContextVar, Token from dataclasses import dataclass from typing import Any -from nemo_gym.token_id_capture.protocols import TokenSink +from nemo_gym.token_id_capture.lineage import assistant_fingerprint, stamp_continuation +from nemo_gym.token_id_capture.protocols import LineageResolution, LineageStore, TokenSink from nemo_gym.token_id_capture.records import ( + ParentResolutionStatus, TokenEntry, extract_token_fields, response_to_output_items, + stamp_lineage, strip_token_fields, ) @@ -51,6 +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. """ rollout_id: str @@ -58,12 +64,44 @@ class CaptureContext: # ``None`` means another process owns record staging. # The context still carries the capture identity. token_sink: TokenSink | None + lineage_store: LineageStore | None = None model: str = "" # ``commit_entry`` sets this after another capture path records the call. committed: bool = False + # Resolve the parent once before dispatch. + # Downstream inference and capture share this immutable decision. + parent_resolution: LineageResolution | None = None + + @property + def parent_call_id(self) -> str | None: + match = self.parent_resolution.match if self.parent_resolution is not None else None + return match.model_call_id if match is not None else None + + @property + def parent_tokens(self) -> list[int]: + match = self.parent_resolution.match if self.parent_resolution is not None else None + return list(match.cumulative_token_ids) if match is not None else [] _CAPTURE_CONTEXT: ContextVar[CaptureContext | None] = ContextVar("nemo_gym_capture_context", default=None) +_STATS_LOCK = threading.Lock() +_RESOLUTION_COUNTS = {"root": 0, "resolved": 0, "unresolved": 0} +_CAPTURE_FAILURES = 0 +_RESOLVER_UNAVAILABLE_NOTED = False + + +def _count_resolution(status_value: str) -> None: + with _STATS_LOCK: + _RESOLUTION_COUNTS[status_value] = _RESOLUTION_COUNTS.get(status_value, 0) + 1 + total = sum(_RESOLUTION_COUNTS.values()) + if total % 1000 == 0: + logger.info("token-capture resolutions: %s", dict(_RESOLUTION_COUNTS)) + + +def capture_health_snapshot() -> dict: + """Return worker-level capture health counters.""" + with _STATS_LOCK: + return {"resolutions": dict(_RESOLUTION_COUNTS), "capture_failures": _CAPTURE_FAILURES} def set_token_sink(context: CaptureContext) -> Token: @@ -89,6 +127,7 @@ async def register_call_intent() -> None: ``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: @@ -98,7 +137,47 @@ async def register_call_intent() -> None: await begin_call(context.rollout_id, context.model_call_id) -async def capture_tokens(response: Any) -> None: +async def resolve_parent(request_messages: list | None) -> None: + """Resolve which recorded call this request continues. + + Use the request representation received from the harness. + 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. + """ + context = _CAPTURE_CONTEXT.get() + if context is None or request_messages is None: + return + try: + if not assistant_fingerprint(request_messages): + context.parent_resolution = LineageResolution(ParentResolutionStatus.ROOT) + elif context.lineage_store is None: + context.parent_resolution = LineageResolution( + ParentResolutionStatus.UNRESOLVED, + reason="resolver_unavailable", + ) + global _RESOLVER_UNAVAILABLE_NOTED + with _STATS_LOCK: + first = not _RESOLVER_UNAVAILABLE_NOTED + _RESOLVER_UNAVAILABLE_NOTED = True + if first: + logger.warning("No lineage resolver is available. Every continuation will be unresolved and masked.") + else: + context.parent_resolution = await context.lineage_store.resolve(context.rollout_id, request_messages) + _count_resolution(context.parent_resolution.status.value) + except Exception: + logger.warning("Could not resolve a parent for rollout %s.", context.rollout_id, exc_info=True) + context.parent_resolution = LineageResolution( + ParentResolutionStatus.UNRESOLVED, + reason="lookup_error", + ) + + +async def capture_tokens( + response: Any, + request_messages: list | None = None, +) -> None: """Record a ``TokenEntry`` from a complete model response. Accept a Pydantic model or dictionary. @@ -127,7 +206,16 @@ async def capture_tokens(response: Any) -> None: # Keep content on the output items. # Store token arrays only on the entry. content_items, token_item_index = strip_token_fields(response_to_output_items(payload)) - + # Reuse the parent selected before dispatch. + # Resolve here only when the caller skipped the pre-dispatch step. + if context.parent_resolution is None and request_messages is not None: + await resolve_parent(request_messages) + resolution = context.parent_resolution + if resolution is None: + resolution = LineageResolution( + ParentResolutionStatus.UNRESOLVED, + reason="not_attempted", + ) entry = TokenEntry( rollout_id=context.rollout_id, model_call_id=context.model_call_id, @@ -141,13 +229,19 @@ async def capture_tokens(response: Any) -> None: token_item_index=token_item_index, created_at=time.time(), ) + if request_messages is not None: + stamp_continuation(entry, list(request_messages)) except Exception: await _capture_failed(context, "build") return - await commit_entry(entry) + await commit_entry(entry, parent_resolution=resolution) -async def commit_entry(entry: TokenEntry) -> None: +async def commit_entry( + entry: TokenEntry, + *, + parent_resolution: LineageResolution | None = None, +) -> None: """Durably record a finished entry against the in-flight call. ``capture_tokens`` extracts arrays from a served response. @@ -172,6 +266,20 @@ async def commit_entry(entry: TokenEntry) -> None: context.committed = True return try: + 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 parent decision is persisted with the same sink write. + stamp_lineage( + entry, + resolution.match.model_call_id if resolution.match is not None else None, + parent_resolution=resolution.status, + ) + entry.parent_resolution_reason = resolution.reason or "" await context.token_sink.put(entry) context.committed = True except Exception: @@ -185,6 +293,12 @@ 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 + if failures % 10 == 0: + logger.error("Training-token capture has failed %d times in this worker.", failures) logger.warning( "Training-token capture failed to %s the record for model call %s of rollout %s.", stage, diff --git a/tests/unit_tests/test_rollout_collection.py b/tests/unit_tests/test_rollout_collection.py index 98bd57b310..fdb000bca3 100644 --- a/tests/unit_tests/test_rollout_collection.py +++ b/tests/unit_tests/test_rollout_collection.py @@ -50,10 +50,13 @@ loads_jsonl_line, ) from nemo_gym.token_id_capture import ( + LineageResolution, + ParentResolutionStatus, TokenCaptureSnapshot, TokenCaptureStore, TokenEntry, clear_token_captures_for_rollouts, + stamp_lineage, ) from nemo_gym.token_id_capture.delivery import ( MASK_SAMPLE_KEY, @@ -65,6 +68,19 @@ ) +class _StubLineageStore: + """Satisfy the normal custom-sink contract in collector-only tests.""" + + async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: + return LineageResolution(ParentResolutionStatus.ROOT) + + def is_process_shared(self) -> bool: + return True + + async def close(self) -> None: + pass + + @pytest.fixture def empty_global_config(monkeypatch: pytest.MonkeyPatch) -> MagicMock: get_global_config_dict = MagicMock(return_value={}) @@ -1150,6 +1166,7 @@ async def test_run_from_config_requires_source_before_dispatch( "all_agents": True, "sink": "framework.capture:Sink", "rebuild_response": True, + "lineage_store": f"{__name__}:_StubLineageStore", } }, ) @@ -1209,6 +1226,7 @@ async def close(self): "all_agents": True, "sink": "framework.capture:Sink", "rebuild_response": True, + "lineage_store": f"{__name__}:_StubLineageStore", } }, ) @@ -1950,17 +1968,17 @@ def _record(output: list | None = None) -> dict: @staticmethod def _capture(store: TokenCaptureStore) -> None: - store.append( - TokenEntry( - rollout_id="0-0", - model_call_id="c1", - prompt_token_ids=[1, 2, 3], - generation_token_ids=[4, 5], - generation_log_probs=[-0.1, -0.2], - output_items=[{"type": "message", "role": "assistant", "content": []}], - token_item_index=0, - ) + entry = TokenEntry( + rollout_id="0-0", + model_call_id="c1", + prompt_token_ids=[1, 2, 3], + generation_token_ids=[4, 5], + generation_log_probs=[-0.1, -0.2], + output_items=[{"type": "message", "role": "assistant", "content": []}], + token_item_index=0, ) + stamp_lineage(entry, None, parent_resolution=ParentResolutionStatus.ROOT) + store.append(entry) async def test_rebuilds_a_rollout_that_has_no_token_ids(self, tmp_path: Path) -> None: store = TokenCaptureStore(tmp_path) diff --git a/tests/unit_tests/test_token_capture_conformance.py b/tests/unit_tests/test_token_capture_conformance.py new file mode 100644 index 0000000000..c236f7ee0f --- /dev/null +++ b/tests/unit_tests/test_token_capture_conformance.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run the conformance kit against Gym's own backends. + +The kit is what an external framework (e.g. NeMo-RL over TransferQueue) runs +against its sink/source/lineage adapters; Gym's file store and an in-memory +backend must both pass every applicable check. +""" + +import asyncio + +import pytest + +from nemo_gym.token_id_capture import ( + FileLineageStore, + InMemoryLineageStore, + TokenCaptureSnapshot, + TokenCaptureStore, + TokenEntry, +) +from nemo_gym.token_id_capture.conformance import run_conformance + + +def test_file_store_passes_all_checks(tmp_path): + passed = asyncio.run( + run_conformance( + lambda: TokenCaptureStore(tmp_path), + lambda: TokenCaptureStore(tmp_path), + lambda: FileLineageStore(tmp_path), + ) + ) + assert "begin_call_custody" in passed + assert "lineage_visibility" in passed + assert len(passed) >= 10 + + +class _MemoryBackend: + """The minimal external-transport shape from the integration tests.""" + + def __init__(self): + self.entries: dict[str, dict[str, TokenEntry]] = {} + self.incomplete: set[str] = set() + self.frozen: dict[str, tuple[str, int]] = {} + self.versions: dict[str, int] = {} + self.lineage = InMemoryLineageStore() + + +class _MemorySink: + def __init__(self, backend): + self.backend = backend + + async def put(self, entry: TokenEntry) -> None: + backend = self.backend + if entry.rollout_id in backend.frozen: + backend.versions[entry.rollout_id] = backend.versions.get(entry.rollout_id, 0) + 1 + raise RuntimeError("frozen") + rollout = backend.entries.setdefault(entry.rollout_id, {}) + previous = rollout.get(entry.model_call_id) + if previous is not None: + if previous != entry: + backend.incomplete.add(entry.rollout_id) + backend.versions[entry.rollout_id] = backend.versions.get(entry.rollout_id, 0) + 1 + raise ValueError("conflicting payload") + return + rollout[entry.model_call_id] = entry + backend.versions[entry.rollout_id] = backend.versions.get(entry.rollout_id, 0) + 1 + await backend.lineage.put(entry) + + async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + self.backend.incomplete.add(rollout_id) + self.backend.versions[rollout_id] = self.backend.versions.get(rollout_id, 0) + 1 + + async def close(self) -> None: + pass + + +class _MemorySource: + def __init__(self, backend): + self.backend = backend + + async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: + backend = self.backend + if rollout_id not in backend.frozen: + backend.versions[rollout_id] = backend.versions.get(rollout_id, 0) + 1 + backend.frozen[rollout_id] = (f"snap-{rollout_id}", backend.versions[rollout_id]) + snapshot_id, version = backend.frozen[rollout_id] + return TokenCaptureSnapshot( + rollout_id=rollout_id, + entries=tuple(backend.entries.get(rollout_id, {}).values()), + incomplete=rollout_id in backend.incomplete, + snapshot_id=snapshot_id, + version=backend.versions[rollout_id], + ) + + async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: + backend = self.backend + frozen = backend.frozen.get(rollout_id) + if frozen is None or frozen[0] != snapshot_id or backend.versions.get(rollout_id) != version: + return False + backend.entries.pop(rollout_id, None) + return True + + async def close(self) -> None: + pass + + +class _MemoryLineage: + def __init__(self, backend): + 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 False + + async def close(self) -> None: + pass + + +def test_memory_backend_passes_applicable_checks(): + backend = _MemoryBackend() + passed = asyncio.run( + run_conformance( + lambda: _MemorySink(backend), + lambda: _MemorySource(backend), + lambda: _MemoryLineage(backend), + ) + ) + # No begin_call on this sink: the custody check is skipped, everything else passes. + assert "begin_call_custody" not in passed + assert "lineage_visibility" in passed + assert len(passed) >= 9 + + +class _FeedLineage: + """The intended external-adapter shape: subclass the base, implement two hooks.""" + + def __new__(cls, backend): + from nemo_gym.token_id_capture import IncrementalLineageStore + + class _Impl(IncrementalLineageStore): + def __init__(self, backend): + super().__init__() + self.backend = backend + + def _fetch_new_entries(self, rollout_id, cursor): + entries = list(self.backend.entries.get(rollout_id, {}).values()) + start = cursor or 0 + items = [(entry, entry.model_call_id) for entry in entries[start:]] + return items, len(entries) + + def _load_entry(self, rollout_id, ref): + return self.backend.entries[rollout_id][ref] + + return _Impl(backend) + + +def test_memory_backend_passes_via_the_incremental_base(): + """An adapter built on IncrementalLineageStore passes the kit with ~15 lines of + backend-specific code — the pattern a TransferQueue adapter follows.""" + backend = _MemoryBackend() + passed = asyncio.run( + run_conformance( + lambda: _MemorySink(backend), + lambda: _MemorySource(backend), + lambda: _FeedLineage(backend), + ) + ) + assert "lineage_visibility" in passed + assert "fresh_client_lineage_visibility" in passed + + +def test_kit_rejects_a_broken_backend(tmp_path): + class _Amnesiac(TokenCaptureStore): + def append(self, entry): # drops writes: put acks without durability + return + + from nemo_gym.token_id_capture.conformance import ConformanceError + + with pytest.raises(ConformanceError): + asyncio.run( + run_conformance( + lambda: _Amnesiac(tmp_path), + lambda: TokenCaptureStore(tmp_path), + ) + ) diff --git a/tests/unit_tests/test_token_capture_golden_vectors.py b/tests/unit_tests/test_token_capture_golden_vectors.py new file mode 100644 index 0000000000..de77a1c82c --- /dev/null +++ b/tests/unit_tests/test_token_capture_golden_vectors.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Golden vectors are the cross-repo wire contract for the lineage hashes. + +An external resolver must reproduce these byte-for-byte or every continuation +silently resolves unresolved. If this test fails, either bump the fingerprint or +digest version machinery deliberately, or fix the regression — never regenerate +the vectors silently. +""" + +from nemo_gym.token_id_capture.lineage import assistant_fingerprint, canonicalize_tool_arguments, conversation_digest +from nemo_gym.token_id_capture.records import compute_digest + + +VECTORS = { + "fingerprint": { + "plain": { + "input": [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi there"}, + ], + "expected": "5b20b01e7ca4ec80db6e221c277c02a796f6a2b90e81c7adaf7f1b8fc2c6846e", # pragma: allowlist secret + }, + "chat_tool": { + "input": [ + {"role": "user", "content": "find x"}, + { + "role": "assistant", + "content": "calling", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "search", "arguments": '{"q": "x", "k": 3}'}, + } + ], + }, + ], + "expected": "82955bc805cccaf9d2faa840479ff6c30420d20e42bf2bc1ebef14d679ec3768", # pragma: allowlist secret + }, + "anthropic_tool": { + "input": [ + {"role": "user", "content": "find x"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "calling"}, + {"type": "tool_use", "id": "c1", "name": "search", "input": {"k": 3, "q": "x"}}, + ], + }, + ], + "expected": "82955bc805cccaf9d2faa840479ff6c30420d20e42bf2bc1ebef14d679ec3768", # pragma: allowlist secret + }, + "responses_tool": { + "input": [ + {"role": "user", "content": "find x"}, + {"role": "assistant", "content": "calling"}, + { + "type": "function_call", + "call_id": "c1", + "name": "search", + "arguments": '{"k":3,"q":"x"}', + }, + ], + "expected": "82955bc805cccaf9d2faa840479ff6c30420d20e42bf2bc1ebef14d679ec3768", # pragma: allowlist secret + }, + "empty": {"input": [], "expected": ""}, + }, + "conversation_digest": { + "plain": { + "input": [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi there"}, + ], + "expected": "db10e497fe3d0ee04e81f35f2e3b7857e4c066f5ccaec3c5fcd279f2793ca81b", # pragma: allowlist secret + }, + "multimodal": { + "input": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": {"url": "https://x/y.png"}}, + ], + }, + {"role": "assistant", "content": "seen"}, + ], + "expected": "18d4045a0813dfbe915b43efde91849b483eaf207ba7d130f00473966dc28591", # pragma: allowlist secret + }, + "tool_result": { + "input": [ + {"role": "tool", "tool_call_id": "c1", "content": "42"}, + {"role": "assistant", "content": "done"}, + ], + "expected": "497f57bc226bd826c8d8cb340f7ab535585c0c86123085c958f880c39e83ec3e", # pragma: allowlist secret + }, + "empty": { + "input": [], + "expected": "edd3713968ca572ae44468ff61fabbec572e8de247023892e8bdf0f55aa8cb9f", # pragma: allowlist secret + }, + }, + "canonicalize_tool_arguments": { + "reordered": {"input": '{"k": 3, "q": "x"}', "expected": '{"k":3,"q":"x"}'}, + "unicode": { + "input": '{"s": "caf\u00e9 \\"quoted\\"", "n": 1.5}', + "expected": '{"n":1.5,"s":"caf\u00e9 \\"quoted\\""}', + }, + "none": {"input": None, "expected": ""}, + }, + "compute_digest": { + "empty": { + "input": [], + "expected": "2e222d36d0c6ae1db1ec18b3f96d9a46df5db61b40137bbfbac4cdf0acde9763", # pragma: allowlist secret + }, + "zero": { + "input": [0], + "expected": "0596efeb21b51ec6eb122c9e470703df19a095a7c5ccb96bf87df0ba2447ba59", # pragma: allowlist secret + }, + "small": { + "input": [1, 2, 3], + "expected": "bf99633051449f0f3248b4995a7c8468b9561d7732447341916085f12ff9ff54", # pragma: allowlist secret + }, + "range1000": { + "input": "range(1000)", + "expected": "9b33490a72c5e86cc27fcc0916a08dbb8f1edeb5f108d010dad1b616bd7f0a2b", # pragma: allowlist secret + }, + }, +} + + +def _digest_input(spec): + return list(range(1000)) if spec == "range(1000)" else spec + + +def test_fingerprint_vectors(): + for name, vector in VECTORS["fingerprint"].items(): + assert assistant_fingerprint(vector["input"]) == vector["expected"], name + + +def test_fingerprint_is_identical_across_dialects(): + values = {VECTORS["fingerprint"][k]["expected"] for k in ("chat_tool", "anthropic_tool", "responses_tool")} + assert len(values) == 1 + + +def test_conversation_digest_vectors(): + for name, vector in VECTORS["conversation_digest"].items(): + assert conversation_digest(vector["input"]) == vector["expected"], name + + +def test_tool_argument_canonicalization_vectors(): + for name, vector in VECTORS["canonicalize_tool_arguments"].items(): + assert canonicalize_tool_arguments(vector["input"]) == vector["expected"], name + + +def test_compute_digest_vectors(): + for name, vector in VECTORS["compute_digest"].items(): + assert compute_digest(_digest_input(vector["input"])) == vector["expected"], name diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index 35d628cc9b..ac61bae51b 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -36,12 +36,13 @@ import pytest from fastapi import Body, Request from fastapi.testclient import TestClient -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from nemo_gym.base_responses_api_model import ( BaseResponsesAPIModelConfig, CaptureStore, SimpleResponsesAPIModel, + _request_messages, read_model_call_records, ) from nemo_gym.openai_utils import ( @@ -52,26 +53,49 @@ ) from nemo_gym.server_utils import ServerClient from nemo_gym.token_id_capture import ( + TOKEN_ENTRY_MIN_SCHEMA_VERSION, TOKEN_ENTRY_RECORD_SCHEMA_VERSION, TOKEN_FIELDS, CaptureContext, + LineageResolution, + ParentResolutionStatus, TokenCaptureStore, TokenEntry, TokenIdCaptureConfig, capture_tokens, commit_entry, + compute_digest, + cumulative_tokens, current_capture_context, extract_token_fields, install_token_sink, register_call_intent, reset_token_sink, + resolve_parent, set_token_sink, + stamp_continuation, + stamp_lineage, ) from nemo_gym.token_id_capture.config import token_id_capture_enabled_for_agent +from nemo_gym.token_id_capture.lineage import ( + FileLineageStore, + LineageIndex, + RolloutLineage, + assistant_fingerprint, + conversation_digest, +) from nemo_gym.token_id_capture.protocols import TokenSource from nemo_gym.token_id_capture.store import make_token_store +_ASSISTANT_TURN = { + "role": "assistant", + "content": "checking", + "tool_calls": [{"function": {"name": "search", "arguments": '{"q":"alpha"}'}}], +} + +_MSG_ARGS = {"model": "downstream-model", "max_tokens": 64} + PTOKS = [1, 2, 3] GTOKS = [4, 5] LPS = [-0.1, -0.2] @@ -145,6 +169,24 @@ def test_token_entry_rejects_mismatched_generation_arrays(): ) +def test_token_entry_rejects_inconsistent_parent_resolution(): + common = { + "rollout_id": "r0", + "model_call_id": "c0", + "prompt_token_ids": PTOKS, + "generation_token_ids": GTOKS, + "generation_log_probs": LPS, + } + with pytest.raises(ValidationError, match="requires parent_call_id"): + TokenEntry(**common, parent_resolution=ParentResolutionStatus.RESOLVED) + with pytest.raises(ValidationError, match="cannot carry parent_call_id"): + TokenEntry( + **common, + parent_resolution=ParentResolutionStatus.ROOT, + parent_call_id="parent", + ) + + # --- store -------------------------------------------------------------------- @@ -237,15 +279,15 @@ def test_dangling_call_intent_marks_frozen_snapshot_incomplete(tmp_path): def test_committed_call_satisfies_its_intent(tmp_path): store = TokenCaptureStore(tmp_path) asyncio.run(store.begin_call("complete", "c1")) - store.append( - TokenEntry( - rollout_id="complete", - model_call_id="c1", - prompt_token_ids=PTOKS, - generation_token_ids=GTOKS, - generation_log_probs=LPS, - ) + entry = TokenEntry( + rollout_id="complete", + model_call_id="c1", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, ) + stamp_lineage(entry, None, parent_resolution=ParentResolutionStatus.ROOT) + asyncio.run(store.put(entry)) snapshot = store.freeze_now("complete") @@ -309,15 +351,15 @@ def test_token_store_freeze_is_atomic_and_conditional_drop_is_race_safe(tmp_path def test_token_store_sweeps_only_old_retired_tombstones(tmp_path): store = TokenCaptureStore(tmp_path) for rollout_id in ("old", "recent", "live"): - store.append( - TokenEntry( - rollout_id=rollout_id, - model_call_id=f"{rollout_id}-c1", - prompt_token_ids=PTOKS, - generation_token_ids=GTOKS, - generation_log_probs=LPS, - ) + entry = TokenEntry( + rollout_id=rollout_id, + model_call_id=f"{rollout_id}-c1", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, ) + stamp_lineage(entry, None, parent_resolution=ParentResolutionStatus.ROOT) + store.append(entry) for rollout_id in ("old", "recent"): snapshot = store.freeze_now(rollout_id) assert asyncio.run(store.drop(rollout_id, snapshot_id=snapshot.snapshot_id, version=snapshot.version)) @@ -340,9 +382,11 @@ def test_token_store_recovers_state_lag_from_the_durable_jsonl_tail(tmp_path): generation_token_ids=GTOKS, generation_log_probs=LPS, ) + stamp_lineage(first, None, parent_resolution=ParentResolutionStatus.ROOT) store.append(first) state_after_first = store.state_path_for("lag").read_bytes() - store.append(first.model_copy(update={"model_call_id": "c2"})) + second = first.model_copy(update={"model_call_id": "c2"}) + store.append(second) store.state_path_for("lag").write_bytes(state_after_first) snapshot = store.freeze_now("lag") @@ -355,7 +399,14 @@ def test_token_store_recovers_state_lag_from_the_durable_jsonl_tail(tmp_path): def _block(**kwargs) -> dict: - return {"token_id_capture": {"enabled": True, "rebuild_response": False, **kwargs}} + return { + "token_id_capture": { + "enabled": True, + "rebuild_response": False, + "allow_unresolved_continuations": True, + **kwargs, + } + } def test_config_disabled_needs_no_dir(): @@ -401,7 +452,7 @@ def test_mask_fraction_limit_defaults_off_and_parses(): def test_agent_capture_selection_uses_static_agent_config_or_all_agents(): config = { - "token_id_capture": {"enabled": True, "rebuild_response": False}, + "token_id_capture": {"enabled": True, "rebuild_response": False, "allow_unresolved_continuations": True}, "captured": {"responses_api_agents": {"implementation": {"token_id_capture": True}}}, "ordinary": {"responses_api_agents": {"implementation": {"token_id_capture": False}}}, } @@ -435,7 +486,13 @@ def test_config_rejects_an_unknown_key(): def test_config_accepts_a_sink_without_constructing_the_consumer_source(): config = TokenIdCaptureConfig.model_validate( - {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_ConfiguredSink"}} + { + "token_id_capture": { + "enabled": True, + "sink": f"{__name__}:_ConfiguredSink", + "allow_unresolved_continuations": True, + } + } ) assert config.token_id_capture.sink == f"{__name__}:_ConfiguredSink" @@ -520,9 +577,15 @@ async def chat_completions( return _training_chat_completion() -def _server(global_config_dict) -> SimpleResponsesAPIModel: +def _server(global_config_dict, *, num_workers: int | None = None) -> SimpleResponsesAPIModel: return _CapturingModel( - config=BaseResponsesAPIModelConfig(host="0.0.0.0", port=8099, entrypoint="", name="srv"), + config=BaseResponsesAPIModelConfig( + host="0.0.0.0", + port=8099, + entrypoint="", + name="srv", + num_workers=num_workers, + ), server_client=MagicMock(spec=ServerClient, global_config_dict=global_config_dict), ) @@ -737,7 +800,7 @@ def _external_mode(tmp_path) -> dict: """Capture on, no destination in this process: records are staged elsewhere.""" return { "observability_enabled": False, - "token_id_capture": {"enabled": True, "rebuild_response": False}, + "token_id_capture": {"enabled": True, "rebuild_response": False, "allow_unresolved_continuations": True}, } @@ -903,7 +966,7 @@ def installed_sink(): def test_installed_sink_receives_entries_without_a_capture_dir(installed_sink): """The framework path: capture on, no directory anywhere, records still arrive.""" - config = {"token_id_capture": {"enabled": True, "rebuild_response": False}} + config = {"token_id_capture": {"enabled": True, "rebuild_response": False, "allow_unresolved_continuations": True}} client = TestClient(_server(config).setup_webserver()) resp = client.post("/ng-rollout/task0-sink0/training-token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 @@ -938,30 +1001,29 @@ async def boom(entry): raise RuntimeError("transport down") monkeypatch.setattr(installed_sink, "put", boom) - client = TestClient(_server({"token_id_capture": {"enabled": True, "rebuild_response": False}}).setup_webserver()) + client = TestClient( + _server( + {"token_id_capture": {"enabled": True, "rebuild_response": False, "allow_unresolved_continuations": True}} + ).setup_webserver() + ) resp = client.post("/ng-rollout/task0-sink1/training-token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 # capture never fails the model call assert installed_sink.incomplete == [("task0-sink1", installed_sink.incomplete[0][1])] -def test_a_sink_without_mark_incomplete_is_logged_not_swallowed(caplog): - """The signal cannot be lost quietly: that is the outcome the failure path exists to stop.""" +def test_a_sink_without_mark_incomplete_is_refused_at_install(): + """The signal cannot be lost quietly: a sink that cannot report a lost call is + refused when installed, matching the startup validation of configured sinks.""" class _PutOnlySink: async def put(self, entry): raise RuntimeError("transport down") - install_token_sink(_PutOnlySink()) - try: - client = TestClient( - _server({"token_id_capture": {"enabled": True, "rebuild_response": False}}).setup_webserver() - ) - with caplog.at_level(logging.ERROR): - resp = client.post("/ng-rollout/task0-sink2/training-token-capture/v1/responses", json={"input": "hi"}) - assert resp.status_code == 200 - assert any("does not implement mark_incomplete" in r.message for r in caplog.records) - finally: - install_token_sink(None) + from nemo_gym.token_id_capture import installed_token_sink + + with pytest.raises(TypeError, match="mark_incomplete"): + install_token_sink(_PutOnlySink()) + assert installed_token_sink() is None def test_commit_entry_records_a_call_with_no_token_fields_on_the_response(installed_sink): @@ -983,7 +1045,9 @@ def test_commit_entry_records_a_call_with_no_token_fields_on_the_response(instal finally: reset_token_sink(token) assert len(installed_sink.entries) == 1 - assert installed_sink.entries[0].generation_token_ids == GTOKS + # The commit step stamps lineage even when the caller skips extraction. + assert installed_sink.entries[0].cum_len == len(PTOKS) + len(GTOKS) + assert installed_sink.entries[0].digest def test_records_carry_a_schema_version(): @@ -1014,7 +1078,15 @@ def _bad_entry(**kwargs): with patch("nemo_gym.token_id_capture.sink.TokenEntry", _bad_entry): client = TestClient( - _server({"token_id_capture": {"enabled": True, "rebuild_response": False}}).setup_webserver() + _server( + { + "token_id_capture": { + "enabled": True, + "rebuild_response": False, + "allow_unresolved_continuations": True, + } + } + ).setup_webserver() ) resp = client.post("/ng-rollout/task0-bad0/training-token-capture/v1/responses", json={"input": "hi"}) @@ -1026,6 +1098,593 @@ def _bad_entry(**kwargs): assert entry_ctor is TokenEntry # patch scoped +def test_capture_stamps_cum_len_and_digest(tmp_path): + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + client.post("/ng-rollout/lineage0-roll0/training-token-capture/v1/responses", json={"input": "hi"}) + (entry,) = TokenCaptureStore(tmp_path).read_entries("lineage0-roll0") + assert entry.cum_len == len(PTOKS) + len(GTOKS) + assert entry.digest == compute_digest(PTOKS + GTOKS) + # A missing parent link makes the builder match strict token prefixes. + assert entry.parent_call_id is None + + +def test_digest_round_trip_and_stamp_lineage(): + entry = TokenEntry( + rollout_id="r", + model_call_id="c", + prompt_token_ids=[1, 2], + generation_token_ids=[3], + generation_log_probs=[-0.5], + ) + stamp_lineage(entry, "parent-1") + assert cumulative_tokens(entry) == [1, 2, 3] + assert entry.cum_len == 3 + assert entry.parent_call_id == "parent-1" + assert entry.digest == compute_digest([1, 2, 3]) + # Distinct sequences must not collide. + # The empty sequence has a stable digest. + assert compute_digest([1, 2, 3]) != compute_digest([1, 2, 4]) + assert compute_digest([]) == compute_digest([]) + with pytest.raises(ValueError): + compute_digest([-1]) + + +def test_fingerprint_ignores_non_assistant_turns(): + """Use only model-authored turns for lineage lookup.""" + a = assistant_fingerprint([{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}]) + b = assistant_fingerprint([{"role": "user", "content": "DIFFERENT"}, {"role": "assistant", "content": "a"}]) + assert a == b != "" + # A request without an assistant turn starts a new conversation. + assert assistant_fingerprint([{"role": "user", "content": "q"}]) == "" + + +def test_fingerprint_survives_tool_argument_reserialization(): + """Match tool arguments across equivalent JSON serializations.""" + compact = [ + {"role": "assistant", "content": "", "tool_calls": [{"function": {"name": "f", "arguments": '{"b":1,"a":2}'}}]} + ] + pretty = [ + { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "f", "arguments": '{\n "a": 2,\n "b": 1\n}'}}], + } + ] + assert assistant_fingerprint(compact) == assistant_fingerprint(pretty) + + +def test_lineage_resolves_the_parent_across_a_turn(): + lineage = RolloutLineage() + first_request = [{"role": "user", "content": "hello"}] + lineage.record("call-1", first_request + [{"role": "assistant", "content": "hi"}], [1, 2, 3], "d1") + + # The next request echoes the assistant turn. + second_request = first_request + [{"role": "assistant", "content": "hi"}, {"role": "user", "content": "more"}] + parent = lineage.resolve(second_request) + assert parent.status == ParentResolutionStatus.RESOLVED + assert parent.match is not None and parent.match.model_call_id == "call-1" + assert parent.match.cumulative_token_ids == (1, 2, 3) + + +def test_lineage_record_is_idempotent(): + lineage = RolloutLineage() + messages = [{"role": "assistant", "content": "hi"}] + lineage.record("call-1", messages, [1, 2, 3], "d1") + lineage.record("call-1", messages, [1, 2, 3], "d1") + + parent = lineage.resolve(messages) + assert parent.status == ParentResolutionStatus.RESOLVED + assert parent.match is not None and parent.match.model_call_id == "call-1" + + +def test_lineage_rejects_a_conflicting_call_identity(): + lineage = RolloutLineage() + lineage.record("call-1", [{"role": "assistant", "content": "a"}], [1], "d1") + + with pytest.raises(ValueError, match="conflicting lineage record"): + lineage.record("call-1", [{"role": "assistant", "content": "b"}], [2], "d2") + + +def test_lineage_misses_on_a_rewritten_history(): + """Treat compacted or rewritten model history as unresolved.""" + lineage = RolloutLineage() + lineage.record("call-1", [{"role": "assistant", "content": "hi"}], [1, 2, 3], "d1") + assert ( + lineage.resolve([{"role": "assistant", "content": "a summary of the above"}]).status + == ParentResolutionStatus.UNRESOLVED + ) + + +def test_lineage_refuses_an_ambiguous_parent(): + """Refuse to guess between calls with identical output.""" + lineage = RolloutLineage() + messages = [{"role": "assistant", "content": "same"}] + lineage.record("call-a", messages, [1, 2], "da") + lineage.record("call-b", messages, [3, 4], "db") + assert lineage.resolve(messages).status == ParentResolutionStatus.UNRESOLVED + + +def test_lineage_is_a_tree_so_forks_get_the_parent_not_the_previous_call(): + """Resolve both branches to their shared parent. + + A running cursor would give the second branch the first branch's generation. + Exact prefix supply would then consume the wrong cumulative tokens. + """ + lineage = RolloutLineage() + shared = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "plan"}] + lineage.record("parent", shared, [1, 2, 3], "dp") + lineage.record( + "branch-a", + shared + [{"role": "user", "content": "a"}, {"role": "assistant", "content": "A"}], + [1, 2, 3, 4], + "da", + ) + + # The second branch continues the shared parent. + second = shared + [{"role": "user", "content": "b"}] + parent = lineage.resolve(second) + assert parent.status == ParentResolutionStatus.RESOLVED + assert parent.match is not None and parent.match.model_call_id == "parent" + assert parent.match.cumulative_token_ids == (1, 2, 3) + + +def test_lineage_index_is_bounded(): + """Bound worker-local lineage for abandoned rollouts.""" + index = LineageIndex(max_rollouts=3) + for i in range(10): + index.for_rollout(f"r{i}") + assert len(index) == 3 + + +def _put_shared_file_entry( + root: str, + rollout_id: str = "process-shared", + model_call_id: str = "child-process-call", + response_text: str = "hi", +) -> None: + request = [{"role": "user", "content": "hello"}] + entry = TokenEntry( + rollout_id=rollout_id, + model_call_id=model_call_id, + prompt_token_ids=[1, 2], + generation_token_ids=[3], + generation_log_probs=[-0.1], + output_items=[{"role": "assistant", "content": response_text}], + ) + stamp_lineage(entry, None, parent_resolution=ParentResolutionStatus.ROOT) + stamp_continuation(entry, request) + TokenCaptureStore(root).append(entry) + + +async def test_file_lineage_resolves_across_independent_worker_instances(tmp_path): + reader = FileLineageStore(tmp_path) + request = [{"role": "user", "content": "hello"}] + response = [{"role": "assistant", "content": "hi"}] + _put_shared_file_entry(str(tmp_path), "shared-rollout", "call-1") + + parent = await reader.resolve("shared-rollout", request + response + [{"role": "user", "content": "next"}]) + + assert parent.status == ParentResolutionStatus.RESOLVED + assert parent.match is not None + assert parent.match.model_call_id == "call-1" + assert parent.match.cumulative_token_ids == (1, 2, 3) + + +async def test_file_lineage_cache_is_lru_and_metadata_only(tmp_path): + resolver = FileLineageStore(tmp_path, max_cached_rollouts=2) + continuation = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + {"role": "user", "content": "next"}, + ] + for rollout_id in ("r-a", "r-b", "r-c"): + _put_shared_file_entry(str(tmp_path), rollout_id, f"{rollout_id}-c1") + + await resolver.resolve("r-a", continuation) + await resolver.resolve("r-b", continuation) + await resolver.resolve("r-a", continuation) + await resolver.resolve("r-c", continuation) + + assert "r-a" in resolver._cache + assert "r-b" not in resolver._cache + node = resolver._cache["r-a"][2].by_call_id["r-a-c1"] + assert node.cum_tokens is None + assert node.entry_offset >= 0 + + cold = await resolver.resolve("r-b", continuation) + assert cold.status == ParentResolutionStatus.RESOLVED + assert cold.match is not None + assert cold.match.cumulative_token_ids == (1, 2, 3) + + +def test_file_lineage_uses_bounded_striped_locks(tmp_path): + resolver = FileLineageStore(tmp_path) + + for index in range(10_000): + resolver._rollout_lock(f"r-{index}") + + assert len(resolver._rollout_locks) == 256 + + +async def test_file_lineage_appends_without_rewriting_prior_records(tmp_path): + _put_shared_file_entry(str(tmp_path), "shared-rollout", "call-1", "first") + path = tmp_path / "shared-rollout.tokens.jsonl" + first_payload = path.read_bytes() + first_inode = path.stat().st_ino + + _put_shared_file_entry(str(tmp_path), "shared-rollout", "call-2", "second") + + payload = path.read_bytes() + assert path.stat().st_ino == first_inode + assert payload.startswith(first_payload) + assert len(payload.splitlines()) == 2 + + +async def test_file_lineage_resolves_across_spawned_worker_processes(tmp_path): + context = multiprocessing.get_context("spawn") + process = context.Process(target=_put_shared_file_entry, args=(str(tmp_path),)) + process.start() + process.join(timeout=10) + assert process.exitcode == 0 + + store = FileLineageStore(tmp_path) + parent = await store.resolve( + "process-shared", + [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + {"role": "user", "content": "next"}, + ], + ) + assert parent.status == ParentResolutionStatus.RESOLVED + assert parent.match is not None and parent.match.model_call_id == "child-process-call" + + +async def test_failed_token_commit_never_becomes_resolver_visible(tmp_path, monkeypatch): + token_store = TokenCaptureStore(tmp_path) + resolver = FileLineageStore(tmp_path) + request = [{"role": "user", "content": "hello"}] + entry = TokenEntry( + rollout_id="failed-publication", + model_call_id="call-1", + prompt_token_ids=[1, 2], + generation_token_ids=[3], + generation_log_probs=[-0.1], + output_items=[{"role": "assistant", "content": "hi"}], + ) + stamp_lineage(entry, None, parent_resolution=ParentResolutionStatus.ROOT) + stamp_continuation(entry, request) + + monkeypatch.setattr(token_store, "append", MagicMock(side_effect=RuntimeError("write failed"))) + with pytest.raises(RuntimeError, match="write failed"): + await token_store.put(entry) + + resolution = await resolver.resolve( + entry.rollout_id, + request + entry.output_items + [{"role": "user", "content": "next"}], + ) + assert resolution.status == ParentResolutionStatus.UNRESOLVED + + +async def test_retirement_invalidates_warm_worker_indexes(tmp_path): + _put_shared_file_entry(str(tmp_path), "retired-rollout", "call-1") + resolver = FileLineageStore(tmp_path) + continuation = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + {"role": "user", "content": "next"}, + ] + assert (await resolver.resolve("retired-rollout", continuation)).status == ParentResolutionStatus.RESOLVED + + token_store = TokenCaptureStore(tmp_path) + lock_inode = token_store.lock_path_for("retired-rollout").stat().st_ino + snapshot = await token_store.freeze("retired-rollout") + assert await token_store.drop( + "retired-rollout", + snapshot_id=snapshot.snapshot_id, + version=snapshot.version, + ) + + assert (await resolver.resolve("retired-rollout", continuation)).status == ParentResolutionStatus.UNRESOLVED + assert token_store.lock_path_for("retired-rollout").stat().st_ino == lock_inode + with pytest.raises(RuntimeError, match="retired"): + _put_shared_file_entry(str(tmp_path), "retired-rollout", "late-call") + + +def test_served_calls_link_to_their_parent(tmp_path): + """Record the echoed first call as the second call's parent.""" + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + first = [{"role": "user", "content": "hello"}] + client.post("/ng-rollout/lin0-roll0/training-token-capture/v1/chat/completions", json={"messages": first}) + entries = TokenCaptureStore(tmp_path).read_entries("lin0-roll0") + assert len(entries) == 1 and entries[0].parent_call_id is None + assert entries[0].parent_resolution == ParentResolutionStatus.ROOT + assert entries[0].continuation_fingerprint + assert entries[0].continuation_context_digest + + content = entries[0].output_items[0]["content"] + served_text = content if isinstance(content, str) else content[0]["text"] + second = first + [{"role": "assistant", "content": served_text}, {"role": "user", "content": "more"}] + client.post("/ng-rollout/lin0-roll0/training-token-capture/v1/chat/completions", json={"messages": second}) + + entries = TokenCaptureStore(tmp_path).read_entries("lin0-roll0") + assert len(entries) == 2 + assert entries[1].parent_call_id == entries[0].model_call_id + assert entries[1].parent_resolution == ParentResolutionStatus.RESOLVED + + +def test_separate_model_worker_instances_share_committed_lineage(tmp_path): + config = _both_enabled(tmp_path) + worker_a = TestClient(_server(config, num_workers=2).setup_webserver()) + worker_b = TestClient(_server(config, num_workers=2).setup_webserver()) + first = [{"role": "user", "content": "hello"}] + + worker_a.post("/ng-rollout/two-workers/training-token-capture/v1/chat/completions", json={"messages": first}) + first_entry = TokenCaptureStore(tmp_path).read_entries("two-workers")[0] + content = first_entry.output_items[0]["content"] + served_text = content if isinstance(content, str) else content[0]["text"] + second = first + [{"role": "assistant", "content": served_text}, {"role": "user", "content": "more"}] + worker_b.post("/ng-rollout/two-workers/training-token-capture/v1/chat/completions", json={"messages": second}) + + entries = TokenCaptureStore(tmp_path).read_entries("two-workers") + assert len(entries) == 2 + assert entries[1].parent_resolution == ParentResolutionStatus.RESOLVED + assert entries[1].parent_call_id == entries[0].model_call_id + + +def test_served_calls_do_not_link_across_a_changed_system_prompt(tmp_path): + """Reject lineage across a changed system prompt. + + Anthropic sends the system prompt beside the message list. + Reusing the old prefix would continue instructions that the harness did not send. + """ + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + store = TokenCaptureStore(tmp_path) + + def call(rollout, system, messages): + client.post( + f"/ng-rollout/{rollout}/training-token-capture/v1/messages", + json={"system": system, "messages": messages, **_MSG_ARGS}, + ) + + first = [{"role": "user", "content": "hello"}] + call("sys0", "SYSTEM ONE", first) + served = store.read_entries("sys0")[0] + content = served.output_items[0]["content"] + echoed = content if isinstance(content, str) else content[0]["text"] + second = first + [{"role": "assistant", "content": echoed}, {"role": "user", "content": "more"}] + + # Same conversation, different instructions. + call("sys0", "SYSTEM TWO", second) + entries = store.read_entries("sys0") + assert len(entries) == 2 + assert entries[1].parent_call_id is None + assert entries[1].parent_resolution == ParentResolutionStatus.UNRESOLVED + + # Unchanged instructions preserve the link. + call("sys1", "SYSTEM ONE", first) + call("sys1", "SYSTEM ONE", second) + linked = store.read_entries("sys1") + assert len(linked) == 2 + assert linked[1].parent_call_id == linked[0].model_call_id + assert linked[1].parent_resolution == ParentResolutionStatus.RESOLVED + + +async def test_lineage_lookup_failure_is_persisted_as_unresolved(tmp_path): + class _FailingResolver: + async def resolve(self, rollout_id, request_items): + raise RuntimeError("resolver unavailable") + + def is_process_shared(self): + return True + + async def close(self): + pass + + store = TokenCaptureStore(tmp_path) + context = CaptureContext( + rollout_id="lookup-failure", + model_call_id="call-1", + token_sink=store, + lineage_store=_FailingResolver(), + ) + request = [{"role": "assistant", "content": "prior output"}, {"role": "user", "content": "continue"}] + token = set_token_sink(context) + try: + await resolve_parent(request) + await capture_tokens( + { + "output": [ + { + "type": "message", + "role": "assistant", + "content": "next output", + "prompt_token_ids": PTOKS, + "generation_token_ids": GTOKS, + "generation_log_probs": LPS, + } + ] + }, + request_messages=request, + ) + finally: + reset_token_sink(token) + + entry = store.read_entries("lookup-failure")[0] + assert entry.parent_resolution == ParentResolutionStatus.UNRESOLVED + assert entry.parent_call_id is None + + +def test_a_changed_tool_schema_breaks_the_link(): + """Reject lineage across a changed tool schema.""" + turn = [{"role": "user", "content": "hi"}, _ASSISTANT_TURN] + with_search = _request_messages({"messages": turn, "tools": [{"name": "search"}]}) + with_bash = _request_messages({"messages": turn, "tools": [{"name": "bash"}]}) + + lineage = RolloutLineage() + lineage.record("call-1", with_search, [1, 2, 3], "d1") + assert ( + lineage.resolve(with_search + [{"role": "user", "content": "next"}]).status == ParentResolutionStatus.RESOLVED + ) + assert ( + lineage.resolve(with_bash + [{"role": "user", "content": "next"}]).status == ParentResolutionStatus.UNRESOLVED + ) + + +def test_the_envelope_does_not_change_the_lookup_key(): + """Exclude the request envelope from the assistant fingerprint.""" + turn = [{"role": "user", "content": "hi"}, _ASSISTANT_TURN] + plain = _request_messages({"messages": turn}) + enveloped = _request_messages({"messages": turn, "instructions": "be brief"}) + assert len(enveloped) == len(plain) + 1 + assert assistant_fingerprint(enveloped) == assistant_fingerprint(plain) != "" + + +def test_the_envelope_is_stable_across_dict_and_model_tools(): + """Normalize equivalent dictionary and model tool schemas.""" + + class _Tool(BaseModel): + name: str + + as_dicts = _request_messages({"messages": [], "tools": [{"name": "search"}]}) + as_models = _request_messages({"messages": [], "tools": [_Tool(name="search")]}) + assert as_dicts == as_models + + +def test_fingerprint_matches_across_openai_and_anthropic_tool_shapes(): + """Match equivalent OpenAI and Anthropic tool-call shapes. + + OpenAI records calls in ``tool_calls``. + Anthropic echoes calls in ``tool_use`` content blocks. + """ + recorded = [ + { + "role": "assistant", + "content": "Let me compute that.", + "tool_calls": [{"id": "c1", "function": {"name": "Bash", "arguments": '{"command":"echo 6"}'}}], + } + ] + echoed = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me compute that."}, + {"type": "tool_use", "id": "c1", "name": "Bash", "input": {"command": "echo 6"}}, + ], + } + ] + assert assistant_fingerprint(recorded) == assistant_fingerprint(echoed) != "" + + +def test_fingerprint_agrees_across_all_three_dialects(): + """Hash equivalent turns identically across all dialects. + + Chat puts tool calls on the message. + Anthropic nests tool calls in content blocks. + Responses emits roleless ``function_call`` items. + """ + + anthropic = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "c1", "name": "Bash", "input": {"cmd": "ls"}}]}, + ] + chat = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "function": {"name": "Bash", "arguments": '{"cmd":"ls"}'}}], + }, + ] + responses = [ + {"type": "message", "role": "user", "content": "hi"}, + {"type": "function_call", "name": "Bash", "arguments": '{"cmd":"ls"}', "call_id": "c1"}, + ] + + assert assistant_fingerprint(anthropic) == assistant_fingerprint(chat) == assistant_fingerprint(responses) + assert assistant_fingerprint(responses) != "" + + +def test_responses_tool_calls_are_distinguished(): + """Distinguish Responses turns with different tool arguments.""" + + def turn(cmd): + return [ + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "checking"}]}, + {"type": "function_call", "name": "Bash", "arguments": '{"cmd":"%s"}' % cmd, "call_id": "c1"}, + ] + + assert assistant_fingerprint(turn("ls")) != assistant_fingerprint(turn("rm -rf /")) + + +def test_tool_call_identity_changes_the_fingerprint(): + def turn(call_id): + return [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call_id, + "function": {"name": "Bash", "arguments": '{"cmd":"ls"}'}, + } + ], + } + ] + + assert assistant_fingerprint(turn("call-a")) != assistant_fingerprint(turn("call-b")) + + +def test_multimodal_content_changes_the_conversation_digest(): + def request(url): + return [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "describe"}, + {"type": "input_image", "image_url": url}, + ], + } + ] + + assert conversation_digest(request("https://example/a.png")) != conversation_digest( + request("https://example/b.png") + ) + + +def test_non_object_request_items_fail_closed(): + with pytest.raises(ValueError, match="not an object"): + conversation_digest([{"role": "user", "content": "ok"}, "unsupported"]) + + +def test_lineage_resolves_a_tool_using_turn_echoed_in_anthropic_shape(): + lineage = RolloutLineage() + produced = { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "c1", "function": {"name": "Bash", "arguments": '{"command":"factor 420"}'}}], + } + lineage.record("call-1", [{"role": "user", "content": "factor 420"}, produced], [1, 2, 3], "d1") + + # The harness echoes the turn as Anthropic blocks. + # The harness then appends the tool result. + next_request = [ + {"role": "user", "content": "factor 420"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "c1", "name": "Bash", "input": {"command": "factor 420"}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "420: 2 2 3 5 7"}], + }, + ] + parent = lineage.resolve(next_request) + assert parent.status == ParentResolutionStatus.RESOLVED + assert parent.match is not None and parent.match.model_call_id == "call-1" + assert parent.match.cumulative_token_ids == (1, 2, 3) + + @pytest.mark.parametrize("bad", ["", "a/b", "../escape", "a b"]) def test_an_unsafe_rollout_id_is_rejected(tmp_path, bad): """Reject rollout ids that could escape the store directory.""" @@ -1048,7 +1707,7 @@ def test_a_record_is_readable_as_soon_as_put_returns(tmp_path): generation_log_probs=[-0.1], ) asyncio.run(store.put(entry)) - assert [entry.model_call_id for entry in asyncio.run(store.freeze("r0")).entries] == ["c1"] + assert [e.model_call_id for e in store.read_entries("r0")] == ["c1"] def test_a_rollout_that_lost_a_call_is_distinguishable_from_a_complete_one(tmp_path): @@ -1089,6 +1748,20 @@ async def close(self) -> None: pass +class _ConfiguredLineage: + def __init__(self, namespace: str = "") -> None: + self.namespace = namespace + + async def resolve(self, rollout_id: str, request_items: list[dict]): + return LineageResolution(ParentResolutionStatus.ROOT) + + def is_process_shared(self) -> bool: + return True + + async def close(self) -> None: + pass + + class _NotASink: async def put(self, entry) -> None: pass @@ -1128,6 +1801,7 @@ def test_a_configured_sink_receives_entries(tmp_path): "enabled": True, "rebuild_response": False, "sink": f"{__name__}:_ConfiguredSink", + "allow_unresolved_continuations": True, } } client = TestClient(_server(config).setup_webserver()) @@ -1149,6 +1823,7 @@ def test_a_configured_sink_wins_over_an_installed_one(installed_sink): "enabled": True, "rebuild_response": False, "sink": f"{__name__}:_ConfiguredSink", + "allow_unresolved_continuations": True, } } client = TestClient(_server(config).setup_webserver()) @@ -1171,6 +1846,35 @@ def test_a_sink_receives_its_configured_kwargs(): assert (sink.endpoint, sink.shard) == ("https://dp", 3) +def test_a_lineage_store_receives_its_configured_kwargs(): + config = TokenIdCaptureConfig.model_validate( + _block( + lineage_store=f"{__name__}:_ConfiguredLineage", + lineage_store_kwargs={"namespace": "training-run"}, + ) + ) + lineage = config.build_lineage_store() + assert lineage.namespace == "training-run" + + +def test_multi_worker_custom_capture_requires_shared_lineage(): + config = _block(sink=f"{__name__}:_ConfiguredSink") + with pytest.raises(ValueError, match="process-shared lineage resolver"): + _server(config, num_workers=2).setup_webserver() + + +def test_multi_worker_custom_capture_accepts_configured_shared_lineage(): + config = _block( + sink=f"{__name__}:_ConfiguredSink", + lineage_store=f"{__name__}:_ConfiguredLineage", + ) + _server(config, num_workers=2).setup_webserver() + + +def test_multi_worker_file_capture_uses_process_shared_lineage(tmp_path): + _server(_both_enabled(tmp_path), num_workers=2).setup_webserver() + + def test_a_sink_given_kwargs_it_cannot_take_is_refused_at_startup(): config = TokenIdCaptureConfig.model_validate(_block(sink=f"{__name__}:_KwargSink", sink_kwargs={"nope": 1})) with pytest.raises(ValueError, match="sink_kwargs"): @@ -1185,6 +1889,7 @@ def test_a_sink_that_cannot_report_failures_is_refused_at_startup(): "enabled": True, "rebuild_response": False, "sink": f"{__name__}:_NotASink", + "allow_unresolved_continuations": True, } } ) @@ -1204,6 +1909,7 @@ def test_a_sink_whose_protocol_member_is_not_callable_is_refused(): "enabled": True, "rebuild_response": False, "sink": f"{__name__}:_NotCallableSink", + "allow_unresolved_continuations": True, } } ) @@ -1217,7 +1923,14 @@ def test_a_sink_whose_protocol_member_is_not_callable_is_refused(): ) def test_a_malformed_sink_path_is_refused_at_startup(target, expected): config = TokenIdCaptureConfig.model_validate( - {"token_id_capture": {"enabled": True, "rebuild_response": False, "sink": target}} + { + "token_id_capture": { + "enabled": True, + "rebuild_response": False, + "sink": target, + "allow_unresolved_continuations": True, + } + } ) with pytest.raises(ValueError, match=expected): config.build_sink() @@ -1263,7 +1976,7 @@ def test_the_store_is_a_token_source(tmp_path): generation_log_probs=[-0.1], ) ) - assert [entry.model_call_id for entry in asyncio.run(store.freeze("r0")).entries] == ["c1"] + assert [e.model_call_id for e in store.read_entries("r0")] == ["c1"] # A colocated source can detect a capture failure. # This prevents training on an incomplete rollout. @@ -1283,10 +1996,9 @@ def _entry_fields(**overrides): ) -def test_a_record_older_than_this_reader_is_accepted(): - """Use defaults for fields absent from older records.""" - entry = TokenEntry(**_entry_fields(schema_version=TOKEN_ENTRY_RECORD_SCHEMA_VERSION - 1)) - assert entry.generation_token_ids == [2] +def test_a_record_below_the_schema_floor_is_refused(): + with pytest.raises(ValidationError, match="below the supported minimum"): + TokenEntry(**_entry_fields(schema_version=TOKEN_ENTRY_MIN_SCHEMA_VERSION - 1)) def test_a_record_newer_than_this_reader_is_refused(): @@ -1306,3 +2018,204 @@ def test_a_newer_record_in_the_store_fails_the_read_rather_than_being_skipped(tm with pytest.raises(ValidationError): store.read_entries("r0") + + +def test_digest_and_cum_len_are_filled_for_every_entry(): + """Stamp cumulative length and digest on every entry.""" + empty = TokenEntry( + rollout_id="r", + model_call_id="e", + prompt_token_ids=[], + generation_token_ids=[], + generation_log_probs=[], + ) + stamp_lineage(empty, None) + assert empty.cum_len == 0 and empty.digest == compute_digest([]) + + normal = TokenEntry( + rollout_id="r", + model_call_id="n", + prompt_token_ids=[1, 2], + generation_token_ids=[3], + generation_log_probs=[-0.1], + ) + stamp_lineage(normal, None) + assert normal.cum_len == 3 and normal.digest == compute_digest([1, 2, 3]) + + +def test_a_rewritten_conversation_does_not_resolve_to_the_original_call(): + """Reject an assistant match when the earlier conversation changed.""" + lineage = RolloutLineage() + original = [{"role": "user", "content": "solve task ALPHA"}] + lineage.record("call-1", original + [_ASSISTANT_TURN], cum_tokens=[1, 2, 3], digest="d1") + + compacted = [{"role": "user", "content": "SUMMARY: we were working on task BETA"}, _ASSISTANT_TURN] + + assert lineage.resolve(compacted).status == ParentResolutionStatus.UNRESOLVED + + +def test_appending_a_tool_result_still_resolves(): + """Resolve a continuation after it appends a tool result.""" + lineage = RolloutLineage() + sent = [{"role": "user", "content": "q"}] + lineage.record("call-1", sent + [_ASSISTANT_TURN], cum_tokens=[1, 2, 3], digest="d1") + + continuation = sent + [_ASSISTANT_TURN, {"role": "tool", "content": "search result"}] + + resolved = lineage.resolve(continuation) + assert resolved.status == ParentResolutionStatus.RESOLVED + assert resolved.match is not None and resolved.match.model_call_id == "call-1" + + +def test_two_calls_with_identical_output_resolve_to_neither(): + """Resolve neither call when their outputs are identical.""" + lineage = RolloutLineage() + messages = [{"role": "user", "content": "q"}, _ASSISTANT_TURN] + lineage.record("call-1", messages, cum_tokens=[1, 2], digest="d1") + lineage.record("call-2", messages, cum_tokens=[9, 9], digest="d2") + + assert lineage.resolve(messages).status == ParentResolutionStatus.UNRESOLVED + + +def test_identical_retries_with_the_same_tokens_share_one_parent(): + lineage = RolloutLineage() + messages = [{"role": "user", "content": "q"}, _ASSISTANT_TURN] + lineage.record("call-2", messages, cum_tokens=[1, 2], digest="same") + lineage.record("call-1", messages, cum_tokens=[1, 2], digest="same") + + resolved = lineage.resolve(messages) + + assert resolved.status == ParentResolutionStatus.RESOLVED + assert resolved.match is not None + assert resolved.match.model_call_id == "call-1" + assert resolved.match.cumulative_token_ids == (1, 2) + + +def test_a_conversation_with_no_model_turn_starts_a_new_root(): + """Start a new root when the request has no model-authored turn.""" + lineage = RolloutLineage() + lineage.record("call-1", [{"role": "user", "content": "q"}, _ASSISTANT_TURN], cum_tokens=[1], digest="d") + + assert lineage.resolve([{"role": "user", "content": "a brand new task"}]).status == ParentResolutionStatus.ROOT + assert assistant_fingerprint([{"role": "user", "content": "q"}]) == "" + + +def test_two_forks_of_one_call_both_resolve_to_it(): + """Resolve two forks to the same parent and cumulative tokens.""" + lineage = RolloutLineage() + base = [{"role": "user", "content": "q"}, _ASSISTANT_TURN] + lineage.record("parent", base, cum_tokens=[1, 2, 3], digest="dp") + + a = lineage.resolve(base + [{"role": "tool", "content": "branch A"}]) + b = lineage.resolve(base + [{"role": "tool", "content": "branch B"}]) + + assert a.status == b.status == ParentResolutionStatus.RESOLVED + assert a.match is not None and b.match is not None + assert a.match.model_call_id == b.match.model_call_id == "parent" + assert a.match.cumulative_token_ids == b.match.cumulative_token_ids == (1, 2, 3) + + +def test_recording_a_child_does_not_mutate_its_parent(): + """Keep a parent immutable while recording children.""" + lineage = RolloutLineage() + base = [{"role": "user", "content": "q"}, _ASSISTANT_TURN] + lineage.record("parent", base, cum_tokens=[1, 2, 3], digest="dp") + before = list(lineage.by_call_id["parent"].cum_tokens) + + for i in range(5): + lineage.record(f"child-{i}", base + [{"role": "tool", "content": str(i)}], [7, 7], "dc") + + assert lineage.by_call_id["parent"].cum_tokens == before + + +def test_an_evicted_rollout_resolves_to_nothing_rather_than_to_another_rollout(): + """Resolve an evicted rollout to nothing.""" + index = LineageIndex(max_rollouts=2, max_tokens=10_000_000) + for name in ("r1", "r2", "r3"): + index.for_rollout(name).record(name, [{"role": "user", "content": "q"}, _ASSISTANT_TURN], [1], "d") + + assert ( + index.for_rollout("r1").resolve([{"role": "user", "content": "q"}, _ASSISTANT_TURN]).status + == ParentResolutionStatus.UNRESOLVED + ) + + +def test_the_last_rollout_is_kept_even_over_budget(): + """Keep the only rollout even when it exceeds the token budget.""" + index = LineageIndex(max_rollouts=1, max_tokens=1) + messages = [{"role": "user", "content": "q"}, _ASSISTANT_TURN] + index.for_rollout("r1").record("c1", messages, [1] * 100, "d") + + assert index.for_rollout("r1").resolve(messages).status == ParentResolutionStatus.RESOLVED + + +def test_a_response_echoed_as_several_items_still_resolves(): + """Resolve a response echoed as several items. + + A Responses harness can echo assistant text and a tool call as separate items. + Indexing the served items preserves that shape. + """ + served = [ + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "let me look"}]}, + {"type": "function_call", "name": "search", "arguments": '{"q":"x"}'}, + ] + sent = [{"role": "user", "content": "find x"}] + lineage = RolloutLineage() + lineage.record("call-1", sent + served, cum_tokens=[1, 2, 3], digest="d", context_len=len(sent)) + + continuation = sent + served + [{"type": "function_call_output", "output": "42"}] + + resolved = lineage.resolve(continuation) + assert resolved.status == ParentResolutionStatus.RESOLVED + assert resolved.match is not None and resolved.match.model_call_id == "call-1" + + +def test_reasoning_the_harness_drops_does_not_break_resolution(): + """Ignore standalone reasoning that the harness does not echo.""" + served = [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "answer"}]}] + sent = [{"role": "user", "content": "q"}] + lineage = RolloutLineage() + lineage.record( + "call-1", + sent + [{"type": "reasoning", "summary": [{"type": "summary_text", "text": "thinking"}]}] + served, + cum_tokens=[1, 2], + digest="d", + context_len=len(sent), + ) + + resolved = lineage.resolve(sent + served) + assert resolved.status == ParentResolutionStatus.RESOLVED + assert resolved.match is not None and resolved.match.model_call_id == "call-1" + + +@pytest.mark.parametrize( + "before, after", + [ + # Responses stores the payload under ``output``. + ( + [{"type": "function_call_output", "call_id": "c1", "output": "42 files"}], + [{"type": "function_call_output", "call_id": "c1", "output": "[truncated]"}], + ), + # Anthropic stores the payload under ``content``. + ( + [{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "42 files"}]}], + [{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "[truncated]"}]}], + ), + # Chat stores the payload as plain content. + ( + [{"role": "tool", "tool_call_id": "c1", "content": "42 files"}], + [{"role": "tool", "tool_call_id": "c1", "content": "[truncated]"}], + ), + ], +) +def test_a_rewritten_tool_result_changes_the_conversation_digest(before, after): + """Change the digest when an earlier tool result changes.""" + assert conversation_digest(before) != conversation_digest(after) + + +def test_the_fingerprint_still_ignores_tool_results(): + """Exclude appended tool results from the lookup fingerprint.""" + turn = [{"role": "assistant", "content": "ok"}] + assert assistant_fingerprint(turn) == assistant_fingerprint( + turn + [{"type": "function_call_output", "output": "42 files"}] + ) diff --git a/tests/unit_tests/test_trajectory_builder.py b/tests/unit_tests/test_trajectory_builder.py index 47e3c0df89..8593644d09 100644 --- a/tests/unit_tests/test_trajectory_builder.py +++ b/tests/unit_tests/test_trajectory_builder.py @@ -20,12 +20,14 @@ import pytest from nemo_gym.token_id_capture import ( + ParentResolutionStatus, TokenCaptureSnapshot, assert_prefix_contiguity, - per_request, + compute_digest, prefix_merging, project_chain_to_output_items, project_main_chain_response, + stamp_lineage, token_id_capture_dirs_from_config, trajectories_for_rollout, trajectories_from_source, @@ -34,8 +36,8 @@ from nemo_gym.token_id_capture.store import TokenCaptureStore -def _entry(mcid, prompt, gen, lp=None, created_at=0.0): - return TokenEntry( +def _entry(mcid, prompt, gen, parent=None, lp=None, created_at=0.0): + e = TokenEntry( rollout_id="t0-r0", model_call_id=mcid, model="m", @@ -46,14 +48,19 @@ def _entry(mcid, prompt, gen, 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.) + status = ParentResolutionStatus.RESOLVED if parent is not None else ParentResolutionStatus.ROOT + stamp_lineage(e, parent, parent_resolution=status) + return e # This append-only rollout has three calls. # Each prompt extends the previous prompt and generation. # Interstitial tokens represent tool output or a new user turn. CALL1 = _entry("c1", [1, 2, 3], [10, 11]) -CALL2 = _entry("c2", [1, 2, 3, 10, 11, 4, 5], [12]) -CALL3 = _entry("c3", [1, 2, 3, 10, 11, 4, 5, 12, 6], [13, 14]) +CALL2 = _entry("c2", [1, 2, 3, 10, 11, 4, 5], [12], parent="c1") +CALL3 = _entry("c3", [1, 2, 3, 10, 11, 4, 5, 12, 6], [13, 14], parent="c2") APPEND_ONLY = [CALL1, CALL2, CALL3] @@ -88,7 +95,8 @@ def test_prefix_merging_handles_a_thousand_turns_without_recursive_traversal(): prompt = [1] for turn in range(1_100): generation = [10_000 + turn] - entries.append(_entry(f"call-{turn:04d}", list(prompt), generation)) + parent = f"call-{turn - 1:04d}" if turn else None + entries.append(_entry(f"call-{turn:04d}", list(prompt), generation, parent=parent)) prompt.extend(generation) prompt.append(20_000 + turn) @@ -109,23 +117,6 @@ def test_order_independent(): assert a["output"] == b["output"] -def test_per_request_marks_the_same_generated_tokens(): - # Both builders identify the same sampled tokens. - merged = prefix_merging(APPEND_ONLY) - per_req = per_request(APPEND_ONLY) - assert len(per_req.chains) == 3 - - merged_tokens = _generated_tokens(project_main_chain_response("t0-r0", merged, model="m")) - per_req_tokens = sorted( - tok - for chain in per_req.chains - for item in project_chain_to_output_items(chain) - for tok in (item.get("generation_token_ids") or []) - ) - assert merged_tokens == sorted([10, 11, 12, 13, 14]) - assert per_req_tokens == sorted([10, 11, 12, 13, 14]) - - def test_projection_is_prefix_contiguous(): out = prefix_merging(APPEND_ONLY) response = project_main_chain_response("t0-r0", out, model="m") @@ -180,9 +171,9 @@ def test_contiguity_assert_catches_a_gap(): assert_prefix_contiguity(broken) -def _content_entry(mcid, prompt, gen, text): +def _content_entry(mcid, prompt, gen, text, parent=None): lp = [-0.1] * len(gen) - return TokenEntry( + entry = TokenEntry( rollout_id="t0-r0", model_call_id=mcid, model="m", @@ -200,12 +191,15 @@ def _content_entry(mcid, prompt, gen, text): } ], ) + status = ParentResolutionStatus.RESOLVED if parent is not None else ParentResolutionStatus.ROOT + stamp_lineage(entry, parent, parent_resolution=status) + return entry def test_projection_carries_content_and_stays_contiguous(): entries = [ _content_entry("c1", [1, 2, 3], [10, 11], "first turn"), - _content_entry("c2", [1, 2, 3, 10, 11, 4, 5], [12], "second turn"), + _content_entry("c2", [1, 2, 3, 10, 11, 4, 5], [12], "second turn", parent="c1"), ] out = prefix_merging(entries) resp = project_main_chain_response("t0-r0", out, model="m") @@ -397,30 +391,6 @@ async def close(self): assert "transport unavailable" in built["error"] -def test_single_response_consumer_rejects_per_request_builder(): - class Source: - async def freeze(self, rollout_id): - return TokenCaptureSnapshot( - rollout_id=rollout_id, - entries=(APPEND_ONLY[0],), - incomplete=False, - snapshot_id="snapshot-3", - version=1, - ) - - async def drop(self, rollout_id, *, snapshot_id, version): - return True - - async def close(self): - return None - - built = asyncio.run(trajectories_from_source("t0-r0", Source(), builder="per_request")) - - assert built["mask_sample"] is True - assert built["rebuilt_response"] is None - assert "not supported by single-response delivery" in built["error"] - - def test_consumer_noop_when_disabled_or_absent(tmp_path): assert token_id_capture_dirs_from_config({}) == [] assert trajectories_for_rollout("t0-r0", []) is None @@ -431,34 +401,6 @@ def test_consumer_noop_when_disabled_or_absent(tmp_path): assert missing["rebuilt_response"] is None -def test_ambiguous_parents_are_quarantined(): - # Two roots have identical cumulative sequences. - # A call extends that shared sequence. - # Its parent is ambiguous. - # The builder quarantines the subtree. - a = _entry("a", [1, 2], [7, 8]) - b = _entry("b", [1, 2], [7, 8]) - child = _entry("child", [1, 2, 7, 8, 9], [20]) - out = prefix_merging([a, b, child]) - assert "child" in out.quarantined - # Every emitted chain excludes the quarantined child. - for chain in out.chains: - assert all(link.entry.model_call_id != "child" for link in chain.links) - - -def test_ambiguous_retry_evidence_cannot_collapse_to_an_empty_success(tmp_path): - store = TokenCaptureStore(tmp_path) - for entry in ( - _entry("a", [1, 2], [7, 8]), - _entry("b", [1, 2], [7, 8]), - _entry("child", [1, 2, 7, 8, 9], [20]), - ): - store.append(entry) - built = trajectories_for_rollout("t0-r0", [tmp_path]) - assert built["mask_sample"] is True - assert built["rebuilt_response"] is None - - # --- side calls and chain selection ------------------------------------------- @@ -474,7 +416,9 @@ def test_the_earliest_root_becomes_the_delivered_chain(): # This short auxiliary call starts after the first agent turn completes. side = _entry("side", [9000, 9001], [7, 7, 7], created_at=200.0) real_1 = _entry("real1", list(range(100, 160)), [200, 201, 202, 203], created_at=100.0) - real_2 = _entry("real2", list(range(100, 160)) + [200, 201, 202, 203, 500], [300, 301, 302], created_at=150.0) + real_2 = _entry( + "real2", list(range(100, 160)) + [200, 201, 202, 203, 500], [300, 301, 302], parent="real1", created_at=150.0 + ) out = prefix_merging([side, real_1, real_2]) main = next(c for c in out.chains if c.chain_id == "main") @@ -548,9 +492,10 @@ def test_post_compaction_chain_is_reported_as_dropped(): Metrics report the remaining chain. """ call_1 = _entry("c1", [1, 2, 3], [4, 5]) - call_2 = _entry("c2", [1, 2, 3, 4, 5, 6], [7]) - # The compacted prompt does not extend a captured sequence. + call_2 = _entry("c2", [1, 2, 3, 4, 5, 6], [7], parent="c1") + # A compacted context resolves UNRESOLVED at request time (edited history). call_3 = _entry("c3", [90, 91], [92, 93, 94, 95]) + stamp_lineage(call_3, None, parent_resolution=ParentResolutionStatus.UNRESOLVED) out = prefix_merging([call_1, call_2, call_3]) assert out.notes.chains == 2 @@ -596,7 +541,7 @@ def test_incomplete_capture_masks_the_rollout(tmp_path): def test_clean_rollout_is_not_masked_and_reports_full_delivery(tmp_path): store = TokenCaptureStore(tmp_path) store.append(_entry("c1", [1, 2, 3], [4, 5])) - store.append(_entry("c2", [1, 2, 3, 4, 5, 6], [7])) + store.append(_entry("c2", [1, 2, 3, 4, 5, 6], [7], parent="c1")) built = trajectories_for_rollout("t0-r0", [tmp_path]) assert built["mask_sample"] is False @@ -682,6 +627,131 @@ def counting_run_builder(entries, builder="prefix_merging"): assert built["metrics"]["n_calls"] == 2 +def _with_lineage(entry, parent_call_id=None): + status = ParentResolutionStatus.RESOLVED if parent_call_id is not None else ParentResolutionStatus.ROOT + stamp_lineage(entry, parent_call_id, parent_resolution=status) + return entry + + +def test_recorded_parent_link_resolves_a_final_call_retry_exactly(): + """Use a recorded parent link to resolve retries exactly. + + The siblings share a prompt and differ only in their generation. + Prefix matching cannot identify which generation the harness kept. + The next call's parent link identifies the survivor. + """ + root = _with_lineage(_entry("root", [1, 2], [3])) + kept = _with_lineage(_entry("kept", [1, 2, 3, 4], [5]), parent_call_id="root") + dropped = _with_lineage(_entry("dropped", [1, 2, 3, 4], [9]), parent_call_id="root") + # The next call explicitly continues ``kept``. + nxt = _with_lineage(_entry("next", [1, 2, 3, 4, 5, 6], [7]), parent_call_id="kept") + + out = prefix_merging([root, kept, dropped, nxt]) + main = next(c for c in out.chains if c.chain_id == "main") + assert [link.entry.model_call_id for link in main.links] == ["root", "kept", "next"] + assert "dropped" in out.quarantined + # Exact resolution leaves no retry unresolved. + assert out.notes.unresolved_retries == [] + + +def test_unresolvable_final_retry_is_flagged_not_silently_tie_broken(): + """Report a final-call retry that has no successor.""" + root = _with_lineage(_entry("root", [1, 2], [3])) + a = _with_lineage(_entry("a", [1, 2, 3, 4], [5]), parent_call_id="root") + b = _with_lineage(_entry("b", [1, 2, 3, 4], [9]), parent_call_id="root") + + out = prefix_merging([root, a, b]) + assert sorted(out.notes.unresolved_retries) == ["a", "b"] + + +def test_a_stale_parent_link_becomes_an_incomplete_fragment(): + """Preserve a call without inventing an edge across a digest mismatch.""" + root = _with_lineage(_entry("root", [1, 2], [3])) + child = _entry("child", [1, 2, 3, 4], [5]) + stamp_lineage(child, "root", parent_resolution=ParentResolutionStatus.RESOLVED) + # Simulate a stale record by corrupting the parent digest. + root.digest = compute_digest([42, 42, 42]) + + out = prefix_merging([root, child]) + assert out.notes.parent_link_failures == {"parent_digest_mismatch": 1} + assert out.notes.unresolved_parent_calls == ["child"] + assert sorted([link.entry.model_call_id for chain in out.chains for link in chain.links]) == ["child", "root"] + assert "child" not in out.quarantined + + +def test_a_missing_filtered_parent_recovers_through_prefix_matching(): + """A recorded parent absent from the build (an empty-generation call, filtered) + is not evidence of conflict: prefix matching reattaches the child to a verified + ancestor instead of dropping the chain. Digest MISMATCH stays fatal.""" + root = _with_lineage(_entry("root", [1, 2], [3])) + empty = _with_lineage(_entry("empty", [1, 2, 3, 4], []), parent_call_id="root") + child = _with_lineage(_entry("child", [1, 2, 3, 4, 5], [6]), parent_call_id="empty") + + out = prefix_merging([root, empty, child]) + + assert out.notes.parent_link_failures == {"parent_call_id_missing_recovered": 1} + assert out.notes.unresolved_parent_calls == [] + assert len(out.chains) == 1 + assert "child" not in out.quarantined + + +def test_a_duplicated_snapshot_entry_is_one_call(): + """An at-least-once transport can deliver one entry twice; the duplicate must not + become a phantom second root that masks a healthy rollout.""" + root = _with_lineage(_entry("root", [1, 2], [3])) + child = _with_lineage(_entry("child", [1, 2, 3, 4], [5]), parent_call_id="root") + duplicate = _with_lineage(_entry("child", [1, 2, 3, 4], [5]), parent_call_id="root") + + out = prefix_merging([root, child, duplicate]) + + assert len(out.chains) == 1 + assert out.notes.unresolved_parent_calls == [] + + +def test_unresolved_parent_never_uses_prefix_fallback(): + root = _with_lineage(_entry("root", [1, 2], [3])) + child = _entry("child", [1, 2, 3, 4], [5]) + stamp_lineage(child, None, parent_resolution=ParentResolutionStatus.UNRESOLVED) + + out = prefix_merging([root, child]) + + assert len(out.chains) == 2 + assert out.notes.unresolved_parent_calls == ["child"] + + +def test_explicit_root_never_uses_prefix_fallback(): + first = _with_lineage(_entry("first", [1, 2], [3])) + root = _with_lineage(_entry("root", [1, 2, 3, 4], [5])) + + out = prefix_merging([first, root]) + + assert len(out.chains) == 2 + assert out.notes.unresolved_parent_calls == [] + + +def test_a_recorded_parent_is_verified_not_trusted(): + """Quarantine a recorded parent with unrelated tokens.""" + previous_attempt = _entry("old", [1, 2, 3], [4, 5]) + this_attempt = _entry("new", [90, 91, 92, 93], [94], parent="old") + + out = prefix_merging([previous_attempt, this_attempt]) + + assert "parent_digest_mismatch" in out.notes.parent_link_failures + for chain in out.chains: + assert [link.entry.model_call_id for link in chain.links] != ["old", "new"] + + +def test_a_correct_parent_link_is_used(): + """Use a recorded parent link that passes verification.""" + parent = _entry("p", [1, 2], [3, 4]) + child = _entry("c", [1, 2, 3, 4, 5], [6], parent="p") + + out = prefix_merging([parent, child]) + + assert out.notes.parent_link_failures == {} + assert any([link.entry.model_call_id for link in chain.links] == ["p", "c"] for chain in out.chains) + + def test_a_chain_that_breaks_is_split_and_reported(): """Split calls whose prompts do not extend each other.