nemo_gym.token_id_capture.external_capture

Backend strategies for framework-owned token capture.

Module Contents

Classes

Name	Description
`ExternalCaptureHandler`	Prepare and finalize one backend-specific external capture call.
`MegatronWorkerCaptureHandler`	Commit lineage after a Megatron Inference worker durably stages a canonical delta.
`VLLMWorkerCaptureHandler`	Commit lineage after a vLLM worker durably stages the token delta.
`_BaseExternalCaptureHandler`	Own the lifecycle shared by external capture backends.

Functions

Name	Description
`_strip_capture_transport_fields`	Keep token IDs, logprobs, routes, and coordinates off the agent hop.
`make_external_capture_handler`	Create the external capture strategy selected by typed configuration.

Data

LOGGER

_MEGATRON_TRANSPORT_FIELDS

API

class nemo_gym.token_id_capture.external_capture.ExternalCaptureHandler()

Protocol

Prepare and finalize one backend-specific external capture call.

nemo_gym.token_id_capture.external_capture.ExternalCaptureHandler.finalize_response(
    served_payload: dict[str, typing.Any]
) -> None

async

Commit lineage from the final API representation served to the client.

nemo_gym.token_id_capture.external_capture.ExternalCaptureHandler.prepare_request(
    request_payload: dict[str, typing.Any]
) -> dict[str, typing.Any]

Attach capture instructions to an admitted engine request.

nemo_gym.token_id_capture.external_capture.ExternalCaptureHandler.prepare_response(
    response_payload: dict[str, typing.Any]
) -> None

Retain the worker acknowledgement and remove capture-only response fields.

Model servers must call this for every completion the worker returns, even one whose acknowledgement is missing. It marks the request as having received a worker completion, and finalize_response commits or poisons the call only when that mark is present. Skipping it for a real completion would leave the call merely uncommitted instead of failing closed with worker_response_missing_commit_coordinates. Completions the model server synthesizes itself (the sequential reasoning guard, or a backend context-limit error converted into an empty completion) never pass through here, so they stay uncommitted.

class nemo_gym.token_id_capture.external_capture.MegatronWorkerCaptureHandler()

Bases: _BaseExternalCaptureHandler

Commit lineage after a Megatron Inference worker durably stages a canonical delta.

_BACKEND_LABEL

= 'Megatron'

nemo_gym.token_id_capture.external_capture.MegatronWorkerCaptureHandler._prepare_admitted_request(
    request_payload: dict[str, typing.Any],
    admission: nemo_gym.token_id_capture.staging.records.CaptureAdmission
) -> dict[str, typing.Any]

class nemo_gym.token_id_capture.external_capture.VLLMWorkerCaptureHandler()

Bases: _BaseExternalCaptureHandler

Commit lineage after a vLLM worker durably stages the token delta.

_BACKEND_LABEL

= 'vLLM'

nemo_gym.token_id_capture.external_capture.VLLMWorkerCaptureHandler._prepare_admitted_request(
    request_payload: dict[str, typing.Any],
    admission: nemo_gym.token_id_capture.staging.records.CaptureAdmission
) -> dict[str, typing.Any]

class nemo_gym.token_id_capture.external_capture._BaseExternalCaptureHandler()

Abstract

Own the lifecycle shared by external capture backends.

_BACKEND_LABEL

str

nemo_gym.token_id_capture.external_capture._BaseExternalCaptureHandler._finalize_admitted_response(
    served_payload: dict[str, typing.Any],
    coords_payload: dict[str, typing.Any] | None,
    context: nemo_gym.token_id_capture.sink.CaptureContext,
    ledger: nemo_gym.token_id_capture.protocols.CaptureLedger,
    admission: nemo_gym.token_id_capture.staging.records.CaptureAdmission
) -> None

async

Validate the worker acknowledgement and commit lineage for an admitted response.

This path operates only on shared Gym contracts (CommitCoords, CallRecord, CaptureLedgerCommit); backends differ only in how _prepare_admitted_request asks the engine to stage tokens.

The ordering invariant the external sink requires — a call must not become a lineage parent until its staged record is durable — holds structurally: the worker stages before acknowledging, so the ledger row (which is what makes the call resolvable) is written only after the coordinates arrive. Custody fields were already stripped from the engine response by prepare_response.

nemo_gym.token_id_capture.external_capture._BaseExternalCaptureHandler._prepare_admitted_request(
    request_payload: dict[str, typing.Any],
    admission: nemo_gym.token_id_capture.staging.records.CaptureAdmission
) -> dict[str, typing.Any]

abstract

Attach backend-specific fields after shared admission checks.

nemo_gym.token_id_capture.external_capture._BaseExternalCaptureHandler.finalize_response(
    served_payload: dict[str, typing.Any]
) -> None

async

Validate the retained acknowledgement and record the served response.

served_payload is the JSON form of the response returned to the client, after conversion succeeded. Fingerprints are computed from that representation so the next turn’s echoed history resolves its parent.

nemo_gym.token_id_capture.external_capture._BaseExternalCaptureHandler.prepare_request(
    request_payload: dict[str, typing.Any]
) -> dict[str, typing.Any]

Attach capture instructions to an engine-bound request.

An unadmitted call (UNRESOLVED — already poisoned in the ledger) is forwarded as plain traffic: the backend captures nothing and the completion still serves the agent.

nemo_gym.token_id_capture.external_capture._BaseExternalCaptureHandler.prepare_response(
    response_payload: dict[str, typing.Any]
) -> None

Strip transport fields and retain the acknowledgement until API conversion finishes.

Lineage is published from the final Chat, Responses, or Messages representation (see finalize_response), so the worker coordinates are parked on the request-scoped capture context and the internal engine response is scrubbed of token data immediately.

nemo_gym.token_id_capture.external_capture._strip_capture_transport_fields(
    payload: dict[str, typing.Any]
) -> None

Keep token IDs, logprobs, routes, and coordinates off the agent hop.

nemo_gym.token_id_capture.external_capture.make_external_capture_handler(
    backend: nemo_gym.token_id_capture.config.ExternalStagingBackend
) -> nemo_gym.token_id_capture.external_capture.ExternalCaptureHandler

Create the external capture strategy selected by typed configuration.

nemo_gym.token_id_capture.external_capture.LOGGER = logging.getLogger(__name__)

nemo_gym.token_id_capture.external_capture._MEGATRON_TRANSPORT_FIELDS = ('compact_prompt_token_ids',)