[tinker] Hold 131k concurrent samples incl. long outputs: aiohttp forwarding, proto-once results, cheap access log, undelivered-result fix, load harness

[avigyabb]

Goal: make sure the Tinker API server (FastAPI/uvicorn/SQLite/in-memory store/forwarding client) is not the bottleneck at 131072 concurrent sample requests on the non-colocated megatron path, and give the Trajectory team a fast, GPU-free way to check that. Based on upstream main (345ce86).

Harness

skyrl/benchmarks/load_test_tinker_sampling.py runs the real skyrl.tinker.api app under uvicorn with backend=megatron, colocate_all=false, the engine subprocess stubbed and the router URL seeded into EngineStateDB, a fake vLLM router (barrier or fixed-latency mode), and a multi-process client that behaves like the Tinker SDK (400 concurrent submits, 45s re-polls, retries on reset) and binds several loopback source IPs so it is not capped by one IP's ~28k ephemeral ports.

uv run --extra tinker python skyrl/benchmarks/load_test_tinker_sampling.py \
    --num-requests 131072 --forwarding-max-connections 2048

# realistic engine queueing: 2048 concurrent generations of 5s each
uv run --extra tinker python skyrl/benchmarks/load_test_tinker_sampling.py \
    --num-requests 131072 --forwarding-max-connections 2048 \
    --vllm-mode latency --max-num-seqs 2048 --gen-seconds 5

# against a real server with the public SDK
uv run --extra tinker python skyrl/benchmarks/load_test_tinker_sampling.py \
    --role load --client sdk --url http://HOST:8000 --num-requests 4096

What it found and what this changes

httpx forwarding pool is quadratic in in-flight requests. httpcore's _assign_requests_to_connections rescans every pooled connection for every request; profiled at 27.5M is_idle calls for 512 requests. Per-request CPU: 14ms at 512 in flight, 27ms at 2048, 38ms at 4096. At 2048 in flight the server needed 119s to forward 2048 instant requests. aiohttp is flat at 0.34ms, so SkyRLTrainInferenceForwardingClient now uses aiohttp: TCPConnector(limit=forwarding_inference_max_connections), sock_read=forwarding_inference_timeout_sec, no total deadline so queued requests wait for the engine instead of hitting the old 300s pool timeout.

Access log through RichHandler cost ~1.5ms of event-loop CPU per HTTP request, 63% of server CPU under load. get_uvicorn_log_config now sends uvicorn.access to a plain handler. Server throughput went 337 -> 568 samples/s.

Undelivered results were evicted, so SDK retries got 404 Future not found (the 128x128 failure Chuck hit). retrieve_future started the short retrieved-TTL clock when the response object was built, even for a poll the SDK had already abandoned after its 45s client timeout; uvicorn drops the send silently, the sweeper evicted the result, and the SDK's retry of the same request_id 404'd (fatal in the SDK). Now the clock starts only if request.is_disconnected() is false, and the retrieved TTL is 300s to outlast the SDK's worst-case re-poll gap (45s timeout + 30s backoff, twice). tests/tinker/test_retrieve_future_lost_response.py reproduces the chain under a real uvicorn socket; it fails on main and passes here.

Config. forwarding_inference_timeout_sec default 300s -> 2048s (with unlimited connections a large burst waits inside vLLM's queue and 128x128 bursts exceed 300s there), and a new sample_max_concurrent_requests (default 2000, the SDK's default) served by /api/v1/client/config so operators can raise the SDK's in-flight cap.

Retry semantics. Connect-phase errors and 5xx rejections (TransientInferenceError) are retried once after refreshing the proxy URL. Read/write failures stay final (NovaSky-AI#2118 stance): vLLM may still be executing the request.

Results (2048 outbound connection cap, SDK-style client, 16 vCPU, no GPU)

Run	Completed	Throughput	Server RSS	healthz p99
main, 2048 requests	2048/2048	17/s (e2e 119s)	1.0 GB	stalls
This PR, 131072 barrier	131072/131072	568/s	1.2 GB	0.8s
This PR, 131072 + 5s engine queueing + 45s re-polls (224k abandoned polls, 889 reconnects)	131072/131072, 0 failures	306/s	4.9 GB	9s
This PR, all 131072 held outstanding for 7 min	133k fds	n/a	6.4 GB	3ms during hold

Long-output rollouts (second commit)

With long outputs the per-result payload work becomes the server's main cost. For a 32k-token result (356KB JSON) the forwarding path spent 4ms (orjson decode, pydantic validate, pydantic JSON dump) and the proto path the SDK >= 0.25 uses on retrieve_future spent another 7ms (stdlib json.loads + proto build) inside a single global lock, capping proto delivery near 150 results/s regardless of concurrency. The proto build holds the GIL, so the thread hop bought nothing.

Now the forwarding client decodes the vLLM body once and encodes straight to SampleResponse wire bytes (serialize_sample_output, pinned byte-for-byte to the validated path by tests). The store keeps the proto bytes (8 bytes/token, 26% smaller than the JSON text); retrieve_future passes them through for proto clients and derives JSON lazily and cached for pre-0.25 clients. Pending entries no longer retain the request body. Store TTLs are EngineConfig fields (external_future_retrieved_ttl_sec, external_future_completed_ttl_sec) because retention after delivery is the dominant memory term: roughly completion rate x result size x window.

Run (proto client, 2048 cap, 5s engine queueing, 45s re-polls)	Completed	Throughput	Peak RSS	healthz p99
Before, 512 concurrent x 32k tokens	512/512	73/s (7.5s server CPU)	0.9 GB	1s stalls
After, 512 concurrent x 32k tokens	512/512	167/s (3.5s server CPU)	0.75 GB	0.06s
After, 131072 x 8k tokens (152k re-polls, 3236 reconnects)	131072/131072, 0 failures	267/s	8.9 GB	8s
After, 32768 x 32k tokens	32768/32768, 0 failures	163/s	9.6 GB	5s

Remaining per-result cost is orjson.loads (0.8ms) plus list-to-numpy conversion (1.9ms) for 32k tokens, so one API process tops out around 165 results/s at 32k tokens and ~270/s at 8k. The 5s stalls at 100% CPU cause some keep-alive resets that the SDK retries; on this branch a re-sent asample is a duplicate generation (Chuck's branch dedups by (model_id, sampling_session_id, seq_id)), which is the next thing worth porting.

Notes

• The Tinker SDK re-polls each outstanding future every 45s, so 131072 in-flight SDK futures means ~1750-2900 retrieve_future requests/s, above what one uvicorn process handles (~1200-2000 HTTP req/s; the in-memory store is single-process). Advertising an 8k-16k sample_max_concurrent_requests and letting the SDK queue the rest is the pragmatic setting; the inference engines bound throughput anyway.
• Unlimited outbound connections cannot exceed ~28k (one source IP's ephemeral ports); set --forwarding-inference-max-connections near engine capacity.
• Chuck's Retain burst results and retry transient inference failures j316chuck/SkyRL#19 fixes the 404 by dropping the retrieved-TTL clock and keeping every result for 2048s; this PR keeps the two-clock design so memory stays bounded when results are large. Same changes stacked on his branch: [tinker] Make the API server hold 131k concurrent samples: aiohttp forwarding, cheap access log, load harness j316chuck/SkyRL#20.
• Tests: tests/tinker/ + tests/utils/ CPU suite, 107 passed.

🤖 Generated with Claude Code

[avigyabb]

we don't want to convert tokens, logprobs, stop_reason to python objects

[avigyabb]

Superseded by the split stack (each PR shows only its own diff and retargets to main as the one below merges):

1. [tinker] Route uvicorn's access log to a plain handler instead of Rich #11
2. [tinker] Forward samples with aiohttp instead of httpx #12
3. [tinker] Keep an undelivered sample result alive for the SDK's retry #13
4. [tinker] Survive completion bursts at the socket layer (accept backlog, keep-alive) #14
5. [tinker] Encode forwarded sample results to proto once and serve them as-is #15
6. [tinker] Decode vLLM completion bodies straight into numpy with pysimdjson #16
7. [tinker] Load harness for the API server's sampling path at 131k concurrency #17