nemo_gym.token_id_capture.lineage

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 supported record distinguishes a root, a resolved parent, and an unresolved boundary. The builder uses token-prefix matching only when a verified parent is absent from the frozen snapshot. It never uses prefix matching to cross an unresolved boundary.

A delivered chain contains exactly the tokens the policy emitted over the recorded context. The hashes ignore reasoning and selected items that a harness may omit when it echoes model output. These differences do not change the captured token sequence. Ambiguous matches remain unresolved rather than risking tokens from the wrong call.

Module Contents

Classes

Name	Description
`FileLineageStore`	Resolve lineage from the token JSONL committed by `TokenCaptureStore`.
`InMemoryLineageStore`	Reference resolver for in-process framework backends and tests.
`IncrementalLineageStore`	Base class for lineage resolvers over any committed-entry backend.
`LineageIndex`	Bound worker-local lineage by rollout and cumulative token counts.
`LineageNode`	-
`RolloutLineage`	Keep an append-only per-rollout call index.

Functions

Name	Description
`_canonical_json`	Serialize JSON-compatible prompt content without losing structure.
`_content_of`	Return typed content parts without discarding prompt-shaping blocks.
`_is_assistant_authored`	Return whether the model produced this item.
`_tool_results_of`	Return tool result identities and payloads across dialects.
`_tools_of`	Return tool calls as `(id, name, canonical arguments)` tuples.
`_update_field`	Hash one tagged, length-delimited UTF-8 field.
`assistant_fingerprint`	Fingerprint the model-authored turns of a request, in order.
`canonicalize_tool_arguments`	Normalize a tool call’s arguments for comparison only.
`conversation_digest`	Hash every turn of a conversation, model-authored or not.
`stamp_continuation`	Add compact lookup metadata before the token entry is committed.

Data

FINGERPRINT_VERSION

_CONTEXT_DOMAIN

_FINGERPRINT_DOMAIN

API

class nemo_gym.token_id_capture.lineage.FileLineageStore(
    root: str | pathlib.Path,
    max_cached_rollouts: int = 65536,
    max_cached_tokens: int = 8000000
)

Bases: 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.

_store

= TokenCaptureStore(root)

nemo_gym.token_id_capture.lineage.FileLineageStore._fetch_new_entries(
    rollout_id: str,
    cursor: typing.Any
) -> tuple[list[tuple[nemo_gym.token_id_capture.records.TokenEntry, typing.Any]], typing.Any]

nemo_gym.token_id_capture.lineage.FileLineageStore._load_entries(
    rollout_id: str,
    refs: list[typing.Any]
) -> list[nemo_gym.token_id_capture.records.TokenEntry]

nemo_gym.token_id_capture.lineage.FileLineageStore._load_entry(
    rollout_id: str,
    ref: typing.Any
) -> nemo_gym.token_id_capture.records.TokenEntry

nemo_gym.token_id_capture.lineage.FileLineageStore._read_locked(
    rollout_id: str
)

class nemo_gym.token_id_capture.lineage.InMemoryLineageStore(
    max_rollouts: int = 512,
    max_tokens: int = 8000000
)

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.

index

nemo_gym.token_id_capture.lineage.InMemoryLineageStore.close() -> None

async

nemo_gym.token_id_capture.lineage.InMemoryLineageStore.is_process_shared() -> bool

nemo_gym.token_id_capture.lineage.InMemoryLineageStore.put(
    entry: nemo_gym.token_id_capture.records.TokenEntry
) -> None

async

Publish one committed entry to the worker-local index.

nemo_gym.token_id_capture.lineage.InMemoryLineageStore.resolve(
    rollout_id: str,
    request_items: list[dict]
) -> nemo_gym.token_id_capture.protocols.LineageResolution

async

class nemo_gym.token_id_capture.lineage.IncrementalLineageStore(
    max_cached_rollouts: int = 65536,
    max_cached_tokens: int = 8000000
)

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.

_cache

dict[str, tuple[Any, dict[str, Any], RolloutLineage]] = {}

_cache_guard

= threading.Lock()

_materialized

dict[str, tuple[str, tuple[int, ...]]] = {}

_materialized_tokens

= 0

_rollout_locks

= tuple((threading.Lock()) for _ in (range(256)))

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._cache_put(
    rollout_id: str,
    value: tuple[typing.Any, dict[str, typing.Any], nemo_gym.token_id_capture.lineage.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.

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._cached_materialized(
    rollout_id: str
) -> tuple[str, tuple[int, ...]] | None

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._fetch_new_entries(
    rollout_id: str,
    cursor: typing.Any
) -> tuple[list[tuple[nemo_gym.token_id_capture.records.TokenEntry, typing.Any]], typing.Any]

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._load_entries(
    rollout_id: str,
    refs: list[typing.Any]
) -> list[nemo_gym.token_id_capture.records.TokenEntry]

Load several committed entries.

Backends can override this hook to fetch a parent chain in one operation.

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._load_entry(
    rollout_id: str,
    ref: typing.Any
) -> nemo_gym.token_id_capture.records.TokenEntry

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._materialize(
    rollout_id: str,
    node: nemo_gym.token_id_capture.lineage.LineageNode,
    refs: dict[str, typing.Any],
    lineage: nemo_gym.token_id_capture.lineage.RolloutLineage
) -> tuple[int, ...]

Load one RESOLVED parent’s cumulative tokens from the backend.

Read the chain in one batch and append each token segment once. Digest verification makes stale references fail closed.

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._read_locked(
    rollout_id: str
)

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._refresh(
    rollout_id: str
) -> tuple[dict[str, typing.Any], nemo_gym.token_id_capture.lineage.RolloutLineage]

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._remember_materialized(
    rollout_id: str,
    call_id: str,
    tokens: tuple[int, ...]
) -> None

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._resolve(
    rollout_id: str,
    request_items: list[dict]
) -> nemo_gym.token_id_capture.protocols.LineageResolution

nemo_gym.token_id_capture.lineage.IncrementalLineageStore._rollout_lock(
    rollout_id: str
)

nemo_gym.token_id_capture.lineage.IncrementalLineageStore.close() -> None

async

nemo_gym.token_id_capture.lineage.IncrementalLineageStore.is_process_shared() -> bool

nemo_gym.token_id_capture.lineage.IncrementalLineageStore.resolve(
    rollout_id: str,
    request_items: list[dict]
) -> nemo_gym.token_id_capture.protocols.LineageResolution

async

class nemo_gym.token_id_capture.lineage.LineageIndex(
    max_rollouts: int = 512,
    max_tokens: int = 8000000
)

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.

_rollouts

dict[str, RolloutLineage] = {}

total_tokens

int

nemo_gym.token_id_capture.lineage.LineageIndex.__len__() -> int

nemo_gym.token_id_capture.lineage.LineageIndex._evict() -> None

nemo_gym.token_id_capture.lineage.LineageIndex.clear() -> None

nemo_gym.token_id_capture.lineage.LineageIndex.drop(
    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.

nemo_gym.token_id_capture.lineage.LineageIndex.for_rollout(
    rollout_id: str
) -> nemo_gym.token_id_capture.lineage.RolloutLineage

class nemo_gym.token_id_capture.lineage.LineageNode(
    call_id: str,
    cum_tokens: list[int] | None,
    cum_len: int,
    digest: str,
    entry_offset: int = -1,
    context_len: int = 0,
    context_digest: str = '',
    parent_call_id: str | None = None,
    prompt_is_delta: bool = False
)

Dataclass

call_id

str

context_digest

str = ''

context_len

int = 0

cum_len

int

cum_tokens

list[int] | None

digest

str

entry_offset

int = -1

parent_call_id

str | None = None

prompt_is_delta

bool = False

class nemo_gym.token_id_capture.lineage.RolloutLineage(
    by_fingerprint: dict[str, list[str]] = dict(),
    by_call_id: dict[str, nemo_gym.token_id_capture.lineage.LineageNode] = dict(),
    total_tokens: int = 0
)

Dataclass

Keep an append-only per-rollout call index.

by_call_id

dict[str, LineageNode] = field(default_factory=dict)

by_fingerprint

dict[str, list[str]] = field(default_factory=dict)

total_tokens

int = 0

nemo_gym.token_id_capture.lineage.RolloutLineage._continues(
    node: nemo_gym.token_id_capture.lineage.LineageNode,
    messages: list[dict]
) -> bool

staticmethod

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.

nemo_gym.token_id_capture.lineage.RolloutLineage.add_entry(
    entry: nemo_gym.token_id_capture.records.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.

nemo_gym.token_id_capture.lineage.RolloutLineage.record(
    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.

nemo_gym.token_id_capture.lineage.RolloutLineage.resolve(
    messages: list[dict]
) -> nemo_gym.token_id_capture.protocols.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.

nemo_gym.token_id_capture.lineage.RolloutLineage.resolve_node(
    messages: list[dict]
) -> tuple[nemo_gym.token_id_capture.records.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.

nemo_gym.token_id_capture.lineage._canonical_json(
    value: typing.Any
) -> str

Serialize JSON-compatible prompt content without losing structure.

nemo_gym.token_id_capture.lineage._content_of(
    content: typing.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.

nemo_gym.token_id_capture.lineage._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.

nemo_gym.token_id_capture.lineage._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.

nemo_gym.token_id_capture.lineage._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.

nemo_gym.token_id_capture.lineage._update_field(
    hasher: typing.Any,
    tag: bytes,
    value: str
) -> None

Hash one tagged, length-delimited UTF-8 field.

nemo_gym.token_id_capture.lineage.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.

nemo_gym.token_id_capture.lineage.canonicalize_tool_arguments(
    value: typing.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.

nemo_gym.token_id_capture.lineage.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.

nemo_gym.token_id_capture.lineage.stamp_continuation(
    entry: nemo_gym.token_id_capture.records.TokenEntry,
    request_items: list[dict]
) -> nemo_gym.token_id_capture.records.TokenEntry

Add compact lookup metadata before the token entry is committed.

nemo_gym.token_id_capture.lineage.FINGERPRINT_VERSION = 1

nemo_gym.token_id_capture.lineage._CONTEXT_DOMAIN = b'nemo-gym-lineage-context'

nemo_gym.token_id_capture.lineage._FINGERPRINT_DOMAIN = b'nemo-gym-lineage'