nemo_gym.token_id_capture.store

Store training TokenEntry records by rollout.

Each rollout uses one <rollout_id>.tokens.jsonl file. Evaluation records use a separate file. Every entry line is fsynced before put returns — that is the durability guarantee. The state index is written atomically but fsynced only on lifecycle transitions (freeze, mark, drop); it is reconstructible from the JSONL tail. A per-rollout file lock serializes writers to the same rollout. Different rollouts can write concurrently.

Module Contents

Classes

Name	Description
`TokenCaptureStore`	Durable, rollout-keyed JSONL sink for `TokenEntry` records.

Functions

Name	Description
`make_token_store`	Build the training-token file store.
`validate_rollout_id`	Reject anything that could escape the store directory or index a bad file.
`validate_rollout_ids`	Validate a whole batch before anything changes, and drop duplicates.

Data

logger

API

class nemo_gym.token_id_capture.TokenCaptureStore(
    root: str | pathlib.Path
)

Durable, rollout-keyed JSONL sink for TokenEntry records.

_root

= Path(root)

root

Path

nemo_gym.token_id_capture.TokenCaptureStore._begin_call(
    rollout_id: str,
    model_call_id: str
) -> None

Durably record that a captured call is about to be dispatched.

A lost entry leaves a dangling intent. freeze_now then masks the rollout. A failure here happens before generation.

nemo_gym.token_id_capture.TokenCaptureStore._dangling_intents(
    rollout_id: str,
    entries: tuple[nemo_gym.token_id_capture.records.TokenEntry, ...]
) -> list[str]

nemo_gym.token_id_capture.TokenCaptureStore._drop(
    rollout_id: str,
    snapshot_id: str,
    version: int
) -> bool

nemo_gym.token_id_capture.TokenCaptureStore._entry_digest(
    payload: bytes
) -> str

staticmethod

nemo_gym.token_id_capture.TokenCaptureStore._fsync_root() -> None

nemo_gym.token_id_capture.TokenCaptureStore._has_records(
    rollout_id: str
) -> bool

nemo_gym.token_id_capture.TokenCaptureStore._locked(
    rollout_id: str,
    shared: bool = False
)

nemo_gym.token_id_capture.TokenCaptureStore._mark_incomplete(
    rollout_id: str,
    model_call_id: str = ''
) -> None

nemo_gym.token_id_capture.TokenCaptureStore._read_entries_unlocked(
    rollout_id: str
) -> list[nemo_gym.token_id_capture.records.TokenEntry]

nemo_gym.token_id_capture.TokenCaptureStore._read_state(
    rollout_id: str
) -> dict[str, typing.Any]

nemo_gym.token_id_capture.TokenCaptureStore._remove_record_files(
    rollout_id: str
) -> None

Remove a rollout’s records, intents, and incomplete marker, whichever exist.

nemo_gym.token_id_capture.TokenCaptureStore._sync_after_error() -> None

nemo_gym.token_id_capture.TokenCaptureStore._sync_entry_index(
    rollout_id: str,
    state: dict[str, typing.Any]
) -> bool

Reconcile an entry index with any durable JSONL tail.

The JSONL write is durable before its state update. A process can therefore stop with one unindexed entry. Normal writes use the state index without parsing prior token arrays. Recovery parses only the unindexed tail.

nemo_gym.token_id_capture.TokenCaptureStore._write_fence_unlocked(
    rollout_id: str,
    state: dict[str, typing.Any]
) -> None

nemo_gym.token_id_capture.TokenCaptureStore._write_state(
    rollout_id: str,
    state: dict[str, typing.Any],
    durable: bool = True,
    sync_root: bool = True
) -> None

nemo_gym.token_id_capture.TokenCaptureStore.append(
    entry: nemo_gym.token_id_capture.records.TokenEntry
) -> None

Idempotently append one entry and fsync.

nemo_gym.token_id_capture.TokenCaptureStore.begin_call(
    rollout_id: str,
    model_call_id: str
) -> None

async

nemo_gym.token_id_capture.TokenCaptureStore.close() -> None

async

The file store owns no persistent handles.

nemo_gym.token_id_capture.TokenCaptureStore.delete(
    rollout_ids: collections.abc.Sequence[str]
) -> nemo_gym.token_id_capture.protocols.RolloutRemovalPayload

async

Remove rollouts’ records and fences.

Delete retired rollouts once nothing of those attempts can still write, for example at the end of the run, which removes their fences. Delete a rollout ID before reusing it, so the new execution starts empty. Deleting again is a no-op.

nemo_gym.token_id_capture.TokenCaptureStore.delete_now(
    rollout_ids: collections.abc.Sequence[str]
) -> nemo_gym.token_id_capture.protocols.RolloutRemovalPayload

Synchronous delete.

Every non-empty batch syncs the directory before returning, even when it found nothing to remove: an overlapping delete may have removed the files without syncing yet, and a caller that reuses the rollout ID after this returns must not see the old fence or records come back after a crash.

nemo_gym.token_id_capture.TokenCaptureStore.drop(
    rollout_id: str,
    snapshot_id: str,
    version: int
) -> bool

async

Delete snapshot payloads while retaining its fence.

nemo_gym.token_id_capture.TokenCaptureStore.freeze(
    rollout_id: str
) -> nemo_gym.token_id_capture.protocols.TokenCaptureSnapshot

async

nemo_gym.token_id_capture.TokenCaptureStore.freeze_now(
    rollout_id: str
) -> nemo_gym.token_id_capture.protocols.TokenCaptureSnapshot

Synchronously freeze one rollout and return its stable snapshot.

nemo_gym.token_id_capture.TokenCaptureStore.incomplete_path_for(
    rollout_id: str
) -> pathlib.Path

Sentinel marking that at least one call of this rollout failed to capture.

nemo_gym.token_id_capture.TokenCaptureStore.intents_path_for(
    rollout_id: str
) -> pathlib.Path

Return the durable per-call intent path.

nemo_gym.token_id_capture.TokenCaptureStore.is_incomplete(
    rollout_id: str
) -> bool

nemo_gym.token_id_capture.TokenCaptureStore.lock_path_for(
    rollout_id: str
) -> pathlib.Path

nemo_gym.token_id_capture.TokenCaptureStore.mark_incomplete(
    rollout_id: str,
    model_call_id: str = ''
) -> None

async

Durably record that a call was lost.

nemo_gym.token_id_capture.TokenCaptureStore.path_for(
    rollout_id: str
) -> pathlib.Path

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

async

Store an entry durably without blocking the event loop.

Await the append so later consumers cannot race a partial file.

nemo_gym.token_id_capture.TokenCaptureStore.read_entries(
    rollout_id: str
) -> list[nemo_gym.token_id_capture.records.TokenEntry]

nemo_gym.token_id_capture.TokenCaptureStore.retire(
    rollout_ids: collections.abc.Sequence[str]
) -> nemo_gym.token_id_capture.protocols.RolloutRemovalPayload

async

Remove rollouts’ records and keep a fence, whatever their state.

This is drop without the snapshot check, for rollouts the consumer is done with but won’t drop as a consumed snapshot: masked or failed captures once their evidence is no longer needed, and abandoned attempts. Like drop, it leaves the fence that makes later writes for the rollout fail, including writes from a duplicate execution of the same rollout. Retiring again is a no-op.

nemo_gym.token_id_capture.TokenCaptureStore.retire_now(
    rollout_ids: collections.abc.Sequence[str]
) -> nemo_gym.token_id_capture.protocols.RolloutRemovalPayload

Synchronous retire.

Two phases, each ending in one directory sync: write every missing fence, then remove the records. The first sync runs for every non-empty batch, even one that writes no fence, because the fences it finds may have been written by an overlapping or crashed retire that has not synced yet: a removal must never become durable before its fence, and the call must not return before the fence is durable.

nemo_gym.token_id_capture.TokenCaptureStore.state_path_for(
    rollout_id: str
) -> pathlib.Path

nemo_gym.token_id_capture.make_token_store(
    global_config_dict: typing.Any
) -> nemo_gym.token_id_capture.store.TokenCaptureStore | None

Build the training-token file store.

Return None when capture is disabled. Return None when no directory resolves. Return None when a custom sink owns the records.

nemo_gym.token_id_capture.validate_rollout_id(
    rollout_id: str
) -> str

Reject anything that could escape the store directory or index a bad file.

nemo_gym.token_id_capture.store.validate_rollout_ids(
    rollout_ids: collections.abc.Sequence[str]
) -> list[str]

Validate a whole batch before anything changes, and drop duplicates.

A bare string is rejected: it is a sequence of one-character rollout IDs.

nemo_gym.token_id_capture.store.logger = logging.getLogger(__name__)