Skip to navigation

nemo_gym.server_runtime.request_handling.request_context

View as Markdown

Middleware that gives each request its own copy of the context variables Gym sets.

Module Contents

Classes

NameDescription
RequestContextIsolationMiddlewareStart every request from a clean context, whatever task uvicorn started it from.

Functions

NameDescription
install_request_context_isolationWrap the complete middleware stack, instrumentation included, in context isolation.

API

class nemo_gym.server_runtime.request_handling.request_context.RequestContextIsolationMiddleware(
app: starlette.types.ASGIApp
)

Start every request from a clean context, whatever task uvicorn started it from.

uvicorn starts a pipelined request’s task from inside the previous request’s task (on_response_complete → loop.create_task), so that task inherits the previous request’s contextvars. The OTel ASGI instrumentation, which runs outermost, has the previous request’s server span current at that point; for the new request it sees a current span, skips traceparent extraction, and files the request as an INTERNAL child of the previous one, under the previous rollout’s trace. Pipelining happens whenever a client reuses a kept-alive connection before this server has sent the terminating body message, which Starlette’s BaseHTTPMiddleware sends late.

Installed outside the instrumentor (see install_request_context_isolation), this resets the OTel context for the request’s duration, so the instrumentor extracts the inbound traceparent. It also clears the rollout id. Today that is a guard only: Gym’s BaseHTTPMiddleware layers run the rollout middleware, /run and the handlers in a child task, so a rollout id set there does not reach the task uvicorn copies; the guard keeps that true if those layers are ever replaced by plain ASGI middleware.

nemo_gym.server_runtime.request_handling.request_context.RequestContextIsolationMiddleware.__call__(
scope: starlette.types.Scope,
receive: starlette.types.Receive,
send: starlette.types.Send
) -> None
async
nemo_gym.server_runtime.request_handling.request_context.install_request_context_isolation(
app: fastapi.FastAPI
) -> None

Wrap the complete middleware stack, instrumentation included, in context isolation.

add_middleware is not enough: the FastAPI instrumentor installs its middleware by overriding build_middleware_stack, so it always ends up outside anything added that way. Wrapping the build the same way, after instrumentation, puts isolation outside the instrumentor, which is the only place it works: the instrumentor must see an empty context so it extracts traceparent, and the handlers must still see the server span the instrumentor then makes current. Call it after instrument_app_for_telemetry.