feat(session): multi-process session-server (N workers + sticky-hash router)

[rmfan]

Problem

The session-server (miles/rollout/session/sessions.py) runs as a single uvicorn process per RolloutManager actor. One GIL serialises tokenizer work (merge_tokens, chat-template render), JSON serialisation of multi-MB SGLang responses, and asyncio scheduling for every in-flight request. On a Slurm CPU node a single process saturates near ~50 r/s. Past that, requests queue inside the event loop and tail latency explodes — production run 1709750 showed one session-server PID serving 180k+ events with row-5 p50 = 20 s and p99 = 192 s, while the upstream SGLang generation itself stayed well under a second.

This PR lifts both layers off a single GIL: the session-server backend can be sharded across N processes, and the front-end proxy that fans requests to those backends can itself be sharded across K processes. Session affinity is preserved end-to-end without any shared state between the K proxy workers.

Design

Backend multi-process (--session-server-workers N)

_start_session_server (in miles/ray/rollout.py) spawns N independent SessionServer processes when N > 1. Each backend gets its own port (walked forward from port+1, tracked in a chosen_ports set so sibling spawns can't race onto the same one), its own tokenizer, its own asyncio loop, and a worker_index it stamps onto every session it creates. Children are daemon=True; the parent installs a SIGTERM handler plus atexit hook that does .terminate() → .join(timeout) → .kill() on each child. N = 1 is bit-for-bit unchanged from today.

Router multi-worker (--session-router-workers K)

The user-facing port is held by a small Python uvicorn proxy (session_router.py) that forwards each request to one of the N backends. With K = 1, the original single-process router runs unchanged. A second bench showed this single router process itself caps at ~50 r/s — even with 4 backends behind it that together can serve ~280 r/s when reached directly, the front-end GIL becomes the new bottleneck (see Evidence below).

With K > 1, _start_session_server spawns K independent uvicorn worker processes, each opening its own listening socket on the user-facing port with SO_REUSEPORT set, then handing the bound socket to uvicorn.Server.run(sockets=[sock]). The Linux kernel (≥ 3.9) hash-distributes inbound TCP connections across the K listeners. The router workers share no state and do not coordinate; each is a stateless duplicate of the others. (uvicorn 0.40–0.49 — the versions we ship — has no reuse_port kwarg, so handing it a pre-bound socket is the supported path.)

How stickiness is preserved

Harbor sees one URL. Every session always lands on the same backend, regardless of which router worker accepted the connection. The mechanism:

1. When a backend handles POST /sessions, SessionRegistry.create_session mints an id of the form w{worker_index}-{uuid4().hex}. The backend's index is encoded in the id itself.
2. On every subsequent request, any router worker that receives it parses the w<idx>- prefix and forwards to backend idx. The id is the routing key — there is no shared hash function and no routing table for the K routers to keep in sync.
3. The kernel's SO_REUSEPORT distribution picks a router worker per inbound TCP connection. That choice is independent of session affinity because every router worker applies the same deterministic prefix → backend rule.
4. Requests with no session_id (POST /sessions, /health) are dispatched round-robin via a per-worker itertools.count(). The chosen backend stamps its own index on the returned id, so all follow-up calls route correctly. Per-worker round-robin counters drift but average out at scale.

Net effect: K routers behave as one logical front-end. Same session_id → same backend, every time, no coordination required.

Supporting changes

• asyncio.to_thread wrap on the two synchronous tokenizer calls (prepare_pretokenized, update_pretokenized_state) — keeps the event loop responsive during merge/encode. Benefits N=1 deployments too.
• Streaming request/response pass-through (client.send(stream=True) + StreamingResponse(aiter_raw())) — no router-side buffering or JSON re-encoding.
• HTTP keepalive enabled in the router → backend httpx pool — eliminates per-request TCP setup at peak.
• Session-id format switched from a bare uuid4().hex (N=1, back-compat) to w{idx}-{uuid4().hex} (N>1) — prefix-encoded, no shared hash algorithm.
• Unknown / malformed session_ids fall back to round-robin instead of returning 404 — lets the existing get_or_create_session auto-reseed path actually fire.

Evidence

Headline server-side server_done_p99 (45 s window, SGLang replaced by a constant-0.5 s stub so the table reflects session-server behaviour alone):

in-flight N	wc=1 backend	wc=4 backends
100	0.55 s	0.55 s
200	5.99 s	6.36 s (stub-queue noise)
400	6.03 s	0.74 s (end-to-end Slurm run)
800	11.63 s	8.17 s

End-to-end on a real Slurm CPU node at N=400 in-flight, going from 1 → 4 backends dropped server_p99 9× (6.74 s → 0.74 s). Throughput at the same point did not scale, which surfaced the second cap.

Router-cap finding (wc=4, N=400):

path	throughput
through the single-process router	~51 r/s
bypassing the router via sticky direct-to-backend addressing	~284 r/s

The 5.5× gap is the router's GIL. K > 1 router workers close it; the K-way path is what makes the multi-backend gains reachable end-to-end through the real entry point.

Files touched

• miles/utils/arguments.py — --session-server-workers, --session-router-workers flags.
• miles/ray/rollout.py _start_session_server — N==1 / K==1 paths unchanged; otherwise spawns N backends and K router workers, registers the SIGTERM reaper before the first .start(), deep-copies per-worker args to avoid shared-ref bugs, and handles port exhaustion explicitly.
• miles/rollout/session/sessions.py — asyncio.to_thread wrap; worker_index surfaced in /health.
• miles/rollout/session/linear_trajectory.py — SessionRegistry(worker_index, worker_count) + prefix-encoded session_id.
• miles/rollout/session/session_router.py (new) — proxy with streaming pass-through, prefix-parsing router, round-robin fallback, SO_REUSEPORT socket helper.
• miles/utils/http_utils.py — httpx keepalive disabled on the loadgen side to spread connections across router workers (kernel hash is by 4-tuple).
• tests/fast/router/test_session_uuid_routing.py, tests/fast/router/test_multi_worker_startup.py, tests/fast/router/test_router_multi_worker.py — prefix invariants across N ∈ {1, 2, 4, 8, 16}, /health schema + malformed-id fallback, K=1 vs K>1 launch contract and per-worker instance_id isolation.

Known limitations

• Backend chat-completions still buffers the SGLang response body. The handler does await response.aread() on the upstream body (miles/rollout/session/session_server.py:73) to extract meta_info. The streaming change in this PR covers the wildcard pass-through path only; the chat-completions path remains buffered. Out of scope.
• Tokenizer memory multiplies by N. Qwen-class tokenizers are approximately a few hundred MB resident per backend; N ≤ 8 is comfortable on current Slurm node memory but should be sized per model.
• SO_REUSEPORT semantics differ on macOS. K > 1 is intended for Linux production. The K=1 default keeps dev environments on the original code path; the multi-worker tests patch the uvicorn call surface rather than binding real ports.
• Router needs K ≈ N for full backend saturation (TODO: swap to nginx). Empirically each Python uvicorn router worker caps at ~50–60 r/s. With K = N the router stops being a coalescer and acts as a pure parallelism layer — same process count as if every backend listened on the user port directly. A follow-up should replace the Python router with nginx (template-generated config, started as a child of _start_session_server): per-worker throughput is ~100× higher, K = 1 nginx worker saturates 8+ backends, and the multi-worker router code path can be removed. Not in this PR to keep the change focused on backend multi-process + minimum-viable routing.
• Stale-update guard previously returned phantom 200s. Follow-up commit (66f27f6) now raises SessionStateConflictError (409) when num_assistant changes during the unlocked proxy phase, instead of silently skipping the state update and returning 200. The old behaviour caused cursor-mismatch crashes in compute_samples_from_openai_records downstream (run 1711903: 24 conflict warnings → 2 cursor-mismatch failures with 88/102-token deltas). Callers must treat 409 from /sessions/{id}/v1/chat/completions as a retryable conflict.

Risk / rollout

Default off (--session-server-workers 1 --session-router-workers 1) → bit-for-bit current behaviour. Recommended initial production config: --session-server-workers 4 --session-router-workers 4 (until the nginx replacement lands, set K ≈ N so the router doesn't become the new cap). Roll forward, no migration; flip per-job after one clean run.

[rmfan]

Audit cleanup pass — B1 + H1–H4

Pushed 5 commits addressing the architectural + blocker + high-severity findings from two prior audits:

• ~/agent-dist/docs/sticky-session-routing-research.md — Mechanism A (Stripe-style prefix-encoded session_id)
• ~/agent-dist/docs/pr-31-code-smell-audit.md — B1, H1, H2, H3, H4

Commits

SHA	Audit item	What changed
27851f8c	H1 + smell #1	Replace md5+rejection-sampling routing with prefix-encoded w<idx>-<uuid> session ids. Router parses the prefix; no shared hash algorithm, no 100-try regen cap, no RuntimeError land-mine. Tests rewritten as parse-correctness invariants; bad ids → 404.
fa89d49c	B1	Router now streams request/response bodies end-to-end. Replaces await request.body() + await response.aread() + json.loads/JSONResponse re-encode with client.send(req, stream=True) + StreamingResponse(aiter_raw()). Hop-by-hop headers factored per RFC 7230 §6.1; upstream response closed in finally to avoid leaks on client disconnect. Incidentally fixes M4 (charset hint preservation).
39efe46a	H2	_register_session_server_reaper now runs before the first .start() call, against a mutable tracked_processes list. SIGTERM during the ~90s ready-window no longer orphans backends.
7ed52a71	H3	Deleted the SIGTERM handler chain. Relies on atexit + daemon=True only (per audit recommendation). Avoids the Ray-init-order race and the closure-cycle on double-init.
b0bb0e16	H4	_shallow_copy_args (copy.copy) → _per_worker_args_copy (copy.deepcopy). Drops the dead dataclasses.replace branch — argparse.Namespace is never a dataclass.

Out of scope (follow-up PR)

M1–M6 and L1–L6 from the code-smell audit. Per the task spec, this PR only lands B1 + H1–H4.

Verification

• python -m py_compile clean on every modified file (linear_trajectory.py, session_router.py, rollout.py, test_session_uuid_routing.py, test_multi_worker_startup.py).
• Existing tests/fast/router/test_session_uuid_routing.py::TestRouterAgreement rewritten to test the new parse contract; tests/fast/router/test_multi_worker_startup.py now mints w<idx>- test ids.

🤖 Generated with Claude Code