API ReferenceFull Library ReferenceNemo GymNemo GymToken Id CaptureStaging

nemo_gym.token_id_capture.staging.capture

View as Markdown

Worker-owned, engine-neutral capture with stage-before-response ordering.

Module Contents

Classes

NameDescription
ActiveCallOne admitted call, the policy version present at admission, and its resolved prefix.
CaptureErrorA caller violated the worker capture lifecycle.
CaptureHostMinimal serving-layer seam used by install_capture.
RolloutTokenCaptureBuild, stage, and acknowledge exact per-call token deltas.
StreamingUnsupportedErrorWorker-owned staging currently accepts only complete responses.

Functions

NameDescription
install_captureWire worker-owned staging into one inference serving layer.

Data

LOGGER

API

class nemo_gym.token_id_capture.staging.ActiveCall(
    admission: nemo_gym.token_id_capture.staging.records.CaptureAdmission,
    weight_version: int,
    prefix_token_ids: list[int] = list()
)
Dataclass

One admitted call, the policy version present at admission, and its resolved prefix.

prefix_token_ids is the exact token prefix the engine prompt must begin with. begin_call resolves it from the admission’s inline prefix or from the caller-supplied ids fetched for a staging_chain. Empty for a text root.

admission
CaptureAdmission
completed
bool = field(default=False, init=False)
model_call_id
str
prefix_token_ids
list[int] = field(default_factory=list)
rollout_id
str
weight_version
int
class nemo_gym.token_id_capture.staging.CaptureError()
Exception

Bases: Exception

A caller violated the worker capture lifecycle.

class nemo_gym.token_id_capture.staging.CaptureHost()

Minimal serving-layer seam used by install_capture.

token_capture
RolloutTokenCapture | None = None
nemo_gym.token_id_capture.staging.CaptureHost.install_token_capture(
    capture: nemo_gym.token_id_capture.staging.capture.RolloutTokenCapture
) -> None

Build, stage, and acknowledge exact per-call token deltas.

_completion_lock
= threading.Lock()
adapter
CaptureAdapter | None
nemo_gym.token_id_capture.staging.RolloutTokenCapture._claim_completion(
    call: nemo_gym.token_id_capture.staging.capture.ActiveCall
) -> None
nemo_gym.token_id_capture.staging.RolloutTokenCapture._failed_coords(
    call: nemo_gym.token_id_capture.staging.capture.ActiveCall
) -> nemo_gym.token_id_capture.staging.records.CommitCoords
staticmethod
nemo_gym.token_id_capture.staging.RolloutTokenCapture._resolve_prefix(
    admission: nemo_gym.token_id_capture.staging.records.CaptureAdmission,
    prefix_token_ids: list[int] | None
) -> list[int]
staticmethod
nemo_gym.token_id_capture.staging.RolloutTokenCapture.begin_call(
    admission: nemo_gym.token_id_capture.staging.records.CaptureAdmission,
    prefix_token_ids: list[int] | None = None,
    stream: bool = False,
    weight_version: int | None = None
) -> nemo_gym.token_id_capture.staging.capture.ActiveCall

Admit a typed gate contract and stamp its generation weight version.

prefix_token_ids is the resolved parent prefix for a token_in call. It is required when the admission stores its prefix as a staging_chain: the caller fetches and concatenates those staged deltas and passes the ids here, the same list it hands to CaptureAdapter.enter_prefix. When the admission carries the prefix inline the argument may be omitted; if given it must match. A text root accepts no prefix. Violations are caller bugs and raise CaptureError; they never poison the call.

weight_version is supplied when engine-finished metadata owns the authoritative per-call epoch. Other worker capture paths leave it unset and read the serving worker’s current version provider instead.

nemo_gym.token_id_capture.staging.RolloutTokenCapture.complete_call(
    call: nemo_gym.token_id_capture.staging.capture.ActiveCall,
    prompt_token_ids: list[int],
    generated_token_ids: list[int],
    generated_logprobs: list[float],
    extras: dict[str, typing.Any] | None = None,
    attachments: collections.abc.Mapping[str, typing.Any] | None = None
) -> nemo_gym.token_id_capture.staging.records.CommitCoords

Stage a normalized delta before returning lightweight coordinates.

attachments are opaque framework payloads staged in the same sink write as the record; they are forwarded untouched and never digested.

nemo_gym.token_id_capture.staging.RolloutTokenCapture.complete_call_from_response(
    call: nemo_gym.token_id_capture.staging.capture.ActiveCall,
    response_payload: typing.Any,
    attachments: collections.abc.Mapping[str, typing.Any] | None = None
) -> nemo_gym.token_id_capture.staging.records.CommitCoords

Extract engine-native material and stage it as one atomic lifecycle step.

nemo_gym.token_id_capture.staging.RolloutTokenCapture.fail_call(
    call: nemo_gym.token_id_capture.staging.capture.ActiveCall,
    reason: str
) -> nemo_gym.token_id_capture.staging.records.CommitCoords

Poison one admitted call that failed before durable staging.

class nemo_gym.token_id_capture.staging.StreamingUnsupportedError()

Bases: CaptureError

Worker-owned staging currently accepts only complete responses.

Wire worker-owned staging into one inference serving layer.

nemo_gym.token_id_capture.staging.capture.LOGGER = logging.getLogger(__name__)