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_outcomes() → dict[str, tuple[int, int]]

Hand over per-backend (successes, failures) since the last drain, and reset.

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 the fleet ledger.

Successes travel with them rather than in a separate call, because the ledger needs both halves of the same window to decide anything: a failure count alone cannot tell a shard that failed three times running from one that failed three times among thousands of successes. Draining them apart would let those windows interleave and reintroduce exactly the bug this fixes.

Only backends with something to report appear, so a quiet window drains empty.

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