Skip to navigation

nemo_gym.models.token_capture.delivery

View as Markdown

Build a token-bearing record from one finished rollout.

The finalizer freezes the rollout’s captured model calls. It rebuilds response.output from that frozen snapshot. It does not retire the snapshot. The caller may retire it only after durable handoff. Retirement uses the frozen snapshot_id and version.

The caller provides both the rollout record and its TokenSource. Gym resolves its source from Gym configuration. A training framework may provide a source from its own data plane. Training correlation must preserve /ng-rollout/<id>/training-token-capture.

Existing token ids are the policy’s sampled data. The finalizer leaves a rollout containing any token ids unchanged. Failed or masked builds retain their capture evidence.

Module Contents

Functions

NameDescription
_unusableMask a rollout that needed token ids and could not get them.
capture_build_can_retireWhether a successful build consumed a frozen snapshot.
finalize_rollout_token_captureRebuild one finished rollout record’s response.output from its recorded token ids.
retire_rollout_token_captureRetire a frozen snapshot after durable handoff.
rollout_carries_token_idsWhether this rollout already holds what training needs.

Data

MASK_SAMPLE_KEY

TERMINAL_CALL_KEY

TERMINAL_RESPONSE_ID_KEY

TOKEN_CAPTURE_KEY

_REDUNDANT_CAPTURE_KEY

API

nemo_gym.models.token_capture.delivery._unusable(
result: dict,
error: str,
message: str
) -> dict

Mask a rollout that needed token ids and could not get them.

An unmasked rollout would appear healthy until it reaches the trainer. The record retains the reason for aggregate reporting.

nemo_gym.models.token_capture.delivery.capture_build_can_retire(
built: dict | None
) -> bool

Whether a successful build consumed a frozen snapshot.

nemo_gym.models.token_capture.delivery.finalize_rollout_token_capture(
result: dict,
mask_incomplete_when_attributed: bool = True
) -> dict | None
async

Rebuild one finished rollout record’s response.output from its recorded token ids.

Call this after the harness and verifier finish the record. The function mutates result in place. It replaces response.output. It preserves the reward and all other harness and verifier output, with one exception: when the rebuilt response is incomplete because the delivered chain’s terminal call was cut (length or content_filter), or because the engine refused the request that continued it for the context window, and the harness reported completed or no status, the harness response takes that status and incomplete_details, so a reader of the record sees why the delivered chain stopped. A harness that reported incomplete or failed itself keeps its own verdict. apply_capture_incomplete_verdict states that rule once, for this path and for a framework that rebuilds from staged records.

The function freezes capture records through source. It rebuilds from that frozen snapshot. It never retires the snapshot. A None source means this caller does not rebuild. A rollout that already carries token ids is left unchanged. Its redundant frozen capture remains eligible for retirement after handoff.

The function never raises. Missing or ambiguous tokens cause masking. Failed or masked builds retain their frozen evidence.

Return the build with its rebuilt response, metrics, and optional error. Return None when no source exists. An unusable build has no rebuilt response and sets mask_sample.

mask_incomplete_when_attributed is the token_id_capture.mask_incomplete_when_attributed setting, which decides whether an incomplete snapshot always masks the rollout.

nemo_gym.models.token_capture.delivery.retire_rollout_token_capture(
rollout_id: str,
built: dict | None
) -> bool
async

Retire a frozen snapshot after durable handoff.

The caller owns the durability boundary. Call this only after downstream acceptance or a local fsync. Retirement uses the frozen snapshot_id and version. Failed or masked builds remain as diagnostic evidence.

nemo_gym.models.token_capture.delivery.rollout_carries_token_ids(
result: dict
) -> bool

Whether this rollout already holds what training needs.

Return true when any output item carries generated token ids. These ids are what the policy sampled. They take precedence over a reconstruction that may differ. Partial token coverage must remain visible instead of being overwritten.

nemo_gym.models.token_capture.delivery.MASK_SAMPLE_KEY = 'mask_sample'
nemo_gym.models.token_capture.delivery.TERMINAL_CALL_KEY = '_ng_terminal_model_call_id'
nemo_gym.models.token_capture.delivery.TERMINAL_RESPONSE_ID_KEY = 'terminal_response_id'
nemo_gym.models.token_capture.delivery.TOKEN_CAPTURE_KEY = '_ng_token_capture'
nemo_gym.models.token_capture.delivery._REDUNDANT_CAPTURE_KEY = '_redundant_capture'