[perf][tinker] Resolve awaited requests with one shared batched poller

[avigyabb]

What

retrieve_future waited by polling per caller: every in-flight request opened its own AsyncSession and ran its own SELECT on a 100 ms → 1 s backoff. Database load therefore scaled with concurrency, and because a session checkout comes from a pool the whole API server shares, those pollers starved every other endpoint.

This replaces it with one background task that resolves all waiters with a single batched query per tick.

Independent of the engine-side PRs in this series (#1992) — that one is the engine subprocess, this is the API server — so it can land in any order.

Why it matters more than it looks

The API's async engine takes SQLAlchemy's defaults: 15 connections (5 + 10 overflow) with a 30 s checkout timeout. Nothing configures them. So with a few thousand concurrent rollouts, thousands of pollers contend for 15 slots and anything else needing the database queues behind them.

session_heartbeat is the canary — a trivial single-row UPDATE that should never be slow. Measured at the DB layer (no GPU, no HTTP server), 2048 concurrent waiters:

	per-caller polling	shared poller
polls demanded vs. achieved	2048/s → 1215/s (pool saturated)	1 query per 50 ms tick, flat in waiter count
session_heartbeat p50	778 ms	2 ms
session_heartbeat max	3820 ms	7 ms

This also reproduces a failure reported from a real run. A user on FSDP fully-async RL hit sqlalchemy.exc.TimeoutError: QueuePool limit of size 5 overflow 10 reached, connection timed out, timeout 30.00 alongside HTTP 500s on /asample. Replaying that shape at the DB layer — 8192 concurrent waiters, which is what groups_per_batch=512 × group_size=4 with several batches off-policy produces — gives:

same load, 120 s	main	this PR
QueuePool timeouts → HTTP 500	1	0
worst submit latency	21–29 s (against the 30 s ceiling)	0.37–1.58 s
submits completed	1,090	62,582

At 2048 waiters main peaks around 8 s — bad but under the ceiling; at 8192 it sits at 21–29 s and starts crossing, which is why the report describes intermittent warnings rather than constant failure.

Worth being precise about the causal chain there, because the original report guessed at it: the weight-update pause doesn't reject generation. It stalls forwarded samples, so in-flight requests pile up, so more concurrent pollers exist, and that exhausts the pool — which is exactly why it clusters around weight updates.

How it works

Two module-level functions in api.py over a dict[int, set[asyncio.Future]] held on app.state:

• wait_for_future(waiters, request_id, timeout) registers a future and awaits it, returning (status, result_data) or None on timeout.
• poll_futures(db_engine, waiters, interval) runs as a lifespan-managed task, and every 50 ms issues one WHERE request_id IN (...) AND status IN (COMPLETED, FAILED) covering every waiter.

Because the query is shared, the interval can be tighter than the old backoff while doing far less work — so this also removes up to ~1 s of latency for a request that finished just after a poll.

Details worth reviewing:

• Each caller gets its own future, not one shared per id. Concurrent waiters on one id are routine rather than exotic: the SDK gives up on a retrieve_future call after 45 s and retries the same request_id, while this endpoint holds for up to 300 s, so anything slower than 45 s accumulates overlapping waiters. Per-caller futures mean a caller giving up removes only its own entry, and an abandoned request can't pin an entry in waiters forever and grow the poll query without bound.
• The endpoint's first query does three jobs at once. It fetches status, result_data, so a missing row is the 404, an already-terminal row returns immediately without waiting a tick, and only a pending row falls through to wait_for_future. Query once → return if terminal → otherwise wait.
• The poller survives errors. A failed iteration is logged and the loop continues, so a transient database problem can't permanently wedge every waiter; callers still have their own timeouts. CancelledError is re-raised ahead of the broad handler so shutdown isn't swallowed.
• Shutdown cancels and awaits the task. The handle lives on app.state.future_poller, which also keeps a strong reference so the task can't be garbage-collected mid-execution.
• No id chunking. SQLite has capped bound parameters at 32766 since 3.32 and Postgres at 65535, both far above any plausible in-flight count.
• Timeout behaviour is unchanged — still 408 after 300 s, now via the named RETRIEVE_FUTURE_TIMEOUT_SECONDS.

Risk

This is the riskiest chunk of the series, which is why it is on its own: it adds a background task to the app lifespan and changes how retrieve_future observes completion. The endpoint's external contract is unchanged — same 200/400/404/408/500 responses and the same bodies — and test_api.py's integration tests exercise it end to end against a real server subprocess, which is the main evidence here.

One behaviour note: a 404 for an unknown id is now decided by the endpoint's own lookup, so it is still immediate.

Testing

uv run --isolated --extra dev --extra jax --extra tinker pytest tests/tinker/ --ignore=tests/tinker/skyrl_train

11 tests in tests/tinker/test_futures.py, covering the waiting primitives and the endpoint:

• resolution once a request completes; failed status surfaced distinctly; None on timeout
• an abandoned request leaves no entry behind — the leak the per-caller design exists to prevent
• one waiter giving up does not strand the others, and multiple waiters on one id all resolve
• query count does not scale with waiter count — 50 concurrent waiters over several ticks must issue fewer than 50 statements, which is the whole point and would fail loudly if per-caller polling came back
• the poller survives a failing iteration
• endpoint: a terminal request is served without consulting the poller, a pending one waits, an unknown id 404s

Two caveats on coverage: there is no GPU E2E run behind this (tests/tinker/skyrl_train/ is excluded and needs 4 GPUs), and the QueuePool reproduction above is at the DB layer — real SQLAlchemy engine, real pool, real waiting code, but no HTTP stack, no vLLM, and no CPU contention from a trainer. So the mechanism and its fix are demonstrated; a full fully-async E2E confirmation is still wanted.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces a shared future waiting mechanism (FutureWaiter) to optimize database connection pool usage by batching queries for in-flight requests, replacing the previous per-caller polling approach. It includes a benchmark script, integration into the FastAPI application, and unit tests. The review feedback highlights three key areas for improvement: explicitly cancelling pending waiters during shutdown to prevent hangs, implementing exponential backoff on background poller failures to avoid log flooding, and handling potential formatting errors when parsing request_id to prevent unhandled 500 errors.