feat(vllm_model): optional session-affinity header for upstream routers

[aoshen02]

Summary

New session_affinity_header config on vllm_model: when set, the per-session OpenAI client sends the Gym session id in that header (e.g. X-Session-ID), so an upstream router can key session→replica affinity on it. Default is unset (no behavior change).

Why

vllm_model already keeps a session on one replica for its whole lifetime, and sha256(session_id) % len(clients) does that consistently across uvicorn workers. What it cannot do is take anything else into account: the assignment is uniform and blind to which replica already holds a prefix, how loaded a replica is, or replicas joining and leaving.

Exporting the session id lets an upstream router make that decision with information Gym does not have, without Gym taking on any routing policy of its own.

Measurements

From NVIDIA-NeMo/RL#3663's replay campaign: 610 recorded coding-agent sessions replayed turn by turn (byte-identical request sequence per arm), Nemotron-3-Super-120B-A12B-BF16, TP2 engines, max_num_batched_tokens 8480, greedy decoding with ignore_eos, 0 preemptions everywhere.

8 nodes / 12 engines, max_concurrent_rollouts: 0 (uncapped)

Arm	Wall	Prefix-cache hit rate
cache_aware router (this header)	13:57	91.6%
consistent_hash router (this header)	15:39	91.8%
no router	15:51	91.8%

16 nodes / 28 engines, max_concurrent_rollouts: 112

Arm	Wall	Prefix-cache hit rate
cache_aware router (this header)	18:57	93.5%
consistent_hash router (this header)	20:26	93.5%
no router	22:03	85.6%

Code under test: this PR's commits on Gym upstream 473f446f (2026-07-24), with NVIDIA-NeMo/RL#3663's commits on NeMo-RL upstream daf46ff3 (2026-08-06). One asymmetry, since it is what separates the two no router rows: the 28-engine no router arm ran on 473f446f's per-uvicorn-process counter assignment, while the 12-engine one ran with sha256(session_id) % len(clients) as on current main. The 85.6% is a property of the counter, not of going routerless.

Read honestly: with the stable assignment underneath, the three arms are level on hit rate, so this header buys nothing on cache locality at 12 engines. The remaining signal is cache_aware's 12% wall-clock lead, and it is not a caching effect — 91.6% is the lowest hit rate of the three. The likely mechanism is first-turn placement: with no prefix to match yet, cache_aware falls back to least-loaded, while a hash places blind; the 28-engine per-engine load spread is consistent with that (1.9x vs 2.5x). Single unreplicated run, so: a hypothesis with support, not a result. Rerunning the routed arms is listed as an open cell in the report; we would rather land the config knob and let the data decide than argue from one run.

Full report: https://github.com/aoshen02/RL/blob/feat/gym-router-url/experiments/routing/README.md

Same shape as #2347; opened alongside the NeMo-RL side so the two halves of the router path can be reviewed together.

Testing

• End-to-end in the campaigns above (router-side decision logs confirm one hash key per session, key lifetime ≈ turns per session).

AI assistance was used for this work; the submitting human has reviewed every line.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[ananthsub]

LGTM! Please address the lint failures. then this is good to merge