nemo_rl.models.generation.generation_router

A NeMo-RL-owned HTTP router in front of the vLLM generation fleet.

NeMo-Gym picks a policy endpoint by static round-robin over a list fixed at process start, with no health input and no failover. A dead vLLM endpoint therefore keeps receiving roughly 1/N of new rollouts for the rest of the run, and Gym retries a refused connection in an uncapped loop with no HTTP timeout.

Rather than change Gym, hand it a single URL that NeMo-RL owns. Gym’s VLLMModelConfig.base_url accepts one string, so its round-robin becomes a no-op and the routing decision moves here, next to the fleet health that already knows which shards are serving.

Two properties make this safe to put in Gym’s critical path:

• The URL never changes. The port is reserved once and passed in, so Ray recreating a restarted actor rebinds the same address. Gym is never reconfigured and never has to fail over – which matters because failing over is exactly what it cannot do.

• Every piece of state is built in init. A restarted actor is immediately usable. This is the deliberate inverse of the NemoGym mistake, where the servers were started from a separate _spinup that Ray never re-runs.

Deliberately not a redirect. Handing Gym a 307 would put its socket back on a vLLM endpoint directly, so a backend dying mid-request would drop it into the same uncapped retry loop this exists to avoid.

One thing this trades away: Gym’s selection is sticky round-robin – a session keeps its backend – so per-request least-outstanding gives up prefix-cache affinity across the turns of a multi-turn rollout. That is a real cost, not purely a defect being fixed, and it is worth measuring before enabling this on a multi-turn workload.

Module Contents

Classes

GenerationRouterImpl	Routing logic and HTTP server, split out so it is testable without Ray.
GenerationRouterActor	Ray actor wrapper. Everything it needs is built in __init__.

Data

_SKIPPED_REQUEST_HEADERS
_SKIPPED_RESPONSE_HEADERS
_STREAM_CHUNK_BYTES

API

nemo_rl.models.generation.generation_router._SKIPPED_REQUEST_HEADERS

‘frozenset(…)’

nemo_rl.models.generation.generation_router._SKIPPED_RESPONSE_HEADERS

‘frozenset(…)’

nemo_rl.models.generation.generation_router._STREAM_CHUNK_BYTES

None

class nemo_rl.models.generation.generation_router.GenerationRouterImpl(

*,

backend_urls: list[str],

host: str,

port: int,

backend_timeout_s: float,

connect_timeout_s: float,

no_healthy_backend_status: int,

health_managed: bool = False,

)

Routing logic and HTTP server, split out so it is testable without Ray.

Initialization

base_url() → str

The single URL handed to NeMo-Gym. Stable for the life of the run.

A method rather than a property so the Ray actor can expose it remotely.

set_serving_backends(urls: list[str]) → None

Replace the eligible backend set.

Takes the full set rather than a delta, so a missed update, a reordered one, or a restarted router all converge on the next push instead of needing sequence numbers and replay.

drain_backend_failures() → dict[str, int]

Hand over the per-backend failure counts and reset them.

The router sees failures no liveness probe can – a wedged engine answers is_alive from a healthy worker process. It holds no monitor reference by design, so instead of reporting, it counts, and the controller’s probe tick drains these into GenerationFleetHealth.report_failure.

metrics() → dict[str, float]

_pick_backend() → Optional[str]

Least-outstanding among eligible backends, or None if there are none.

static _target_url(backend: str, path_qs: str) → str

Map an inbound path onto a backend.

Backends are advertised as http://host:port/v1 while inbound paths already carry their own prefix – /v1/chat/completions for most calls, but bare /tokenize because Gym’s create_tokenize strips /v1 first. Stripping the suffix and appending the full path handles both.

async _handle(request: Any) → Any

_on_backend_error(backend: str, error: BaseException) → Any

Answer for a backend that failed, deliberately rather than by accident.

Without this, aiohttp answers instead, and its choice of status decides whether the run survives. A wedged backend trips the client timeout, which aiohttp reports as 504 – and 504 is in NeMo-Gym’s rate-limit retry subset, where _request_with_retry raises its own ceiling on every attempt. That is an unbounded retry loop at backend_timeout_s per turn: exactly the hang _check_status_is_not_retried_by_gym exists to prevent, reintroduced through the error path that validator never covered.

500 instead, because it is in Gym’s bounded retry set: Gym re-sends this one HTTP call, the next _pick_backend lands on a healthy shard, and a multi-turn rollout keeps the turns it had already completed. The no-healthy-backend status (409) would be wrong here – not retried, so it fails the whole rollout and every turn is redone from scratch by the row re-dispatch a layer up.

async _forward(request: Any, backend: str) → Any

build_app() → Any

Build the aiohttp application serving Gym’s endpoint surface.

serve_in_background() → None

Run the HTTP server on a daemon thread with its own event loop.

The socket is bound here, synchronously, before the thread starts. Bound inside the thread instead, a port conflict raises on a daemon thread nobody awaits: the actor stays alive, base_url() is a pure string format so it keeps resolving, and setup’s “fail here rather than inside Gym” guard never notices. Gym is then handed a URL with no listener and retries the refused connection in an uncapped loop – the exact wedge this router exists to prevent. Binding first turns that into a failed actor construction with the port in the traceback.

Same shape as the vLLM workers handing their reserved socket to uvicorn. Restart stays correct: the replacement process rebinds the port its dead predecessor freed.

is_serving() → bool

class nemo_rl.models.generation.generation_router.GenerationRouterActor(**kwargs: Any)

Bases: nemo_rl.models.generation.generation_router.GenerationRouterImpl

Ray actor wrapper. Everything it needs is built in __init__.

max_restarts=-1 is only meaningful because of that: Ray recreates a restarted actor through __init__ alone, so a class that starts its server from a separate method comes back permanently broken.

Initialization