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
`_reject_multimodal_content`	Fail closed when a Megatron capture request carries media or audio parts.
`_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.

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._reject_multimodal_content(
    request_payload: dict[str, typing.Any]
) -> None

Fail closed when a Megatron capture request carries media or audio parts.

The Megatron adapter stages no media geometry, so a multimodal prompt would commit token rows whose lengths disagree with the expanded engine prompt. Until multimodal staging lands, refuse the request rather than train on it.

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',)