nemo_gym.telemetry.endpoints

View as Markdown

Span wrappers for Gym’s endpoint handlers.

Applied in the three SimpleServer subclasses where routes are registered, so every one of Gym’s ~150 servers is instrumented without touching any of them individually.

Why a wrapper and not the FastAPI auto-instrumentation alone: the instrumentor gives one SERVER span per HTTP request, named after the route. That is the right thing for the transport, but it cannot know that /run is a rollout and /verify is a verification, it cannot be switched on per span group, and it cannot attach Gym’s rollout id. These wrappers add the semantic layer on top.

The rollout id comes from nemo_gym.rollout_correlation.current_rollout_id, the ContextVar Gym already sets from RolloutContextMiddleware and the agent’s /run wrapper. It is bridged onto the span rather than replaced: one correlation scheme, now visible from traces, Gym’s own logs, and captured trajectories alike.

Module Contents

Functions

NameDescription
traced_endpointWrap an async FastAPI handler in a span-group-gated span.
traced_rollout_endpointtraced_endpoint for the agent’s /run, plus gym.rollout.duration_ms.
traced_verify_endpointtraced_endpoint for /verify, plus the gym.verify.* metrics.

Data

ROLLOUT_ID_ATTRIBUTE

API

nemo_gym.telemetry.endpoints.traced_endpoint(
group: str,
span_name: str,
handler: typing.Callable,
static_attributes: typing.Optional[dict] = None
) -> typing.Callable

Wrap an async FastAPI handler in a span-group-gated span.

functools.wraps sets __wrapped__, which is what FastAPI’s inspect.signature follows to build the request model — so the route keeps its body type, its validation and its OpenAPI schema. Gym already relies on this for SimpleResponsesAPIAgent.run_with_rollout_context.

Parameters:

group
str

Span group gating this site. Checked at call time, not decoration time: span groups are configured during init_telemetry, long after import.

span_name
str

Span name, e.g. gym.verify.

handler
Callable

The async handler to wrap.

static_attributes
Optional[dict]Defaults to None

Attributes constant for this route, e.g. the server name. Evaluated once at wrap time, not per request.

Returns: Callable

The wrapped handler.

nemo_gym.telemetry.endpoints.traced_rollout_endpoint(
handler: typing.Callable,
static_attributes: typing.Optional[dict] = None
) -> typing.Callable

traced_endpoint for the agent’s /run, plus gym.rollout.duration_ms.

One /run is one rollout, which makes this the span everything else in a rollout hangs off — the model calls and verifications it triggers become its descendants through W3C context propagation.

nemo_gym.telemetry.endpoints.traced_verify_endpoint(
handler: typing.Callable,
static_attributes: typing.Optional[dict] = None
) -> typing.Callable

traced_endpoint for /verify, plus the gym.verify.* metrics.

succeeded records whether the verification call completed, not whether the task passed. Reward and accuracy are experiment telemetry and belong in W&B, not in an application-telemetry metric — see kb/knowledge/concepts/application-vs-experiment-telemetry.md. A verifier that correctly scores an answer as wrong is a success here; a verifier that raises is not.

nemo_gym.telemetry.endpoints.ROLLOUT_ID_ATTRIBUTE = 'nemo.gym.rollout.id'