[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. That makes database load scale with concurrency, and — because a session checkout comes from a pool the whole API server shares — it starves every other endpoint (training loop, weight sync, session liveness, setup/bookkeeping).

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

Independent of #1976 and #1977 (those are engine-side; this is API-side), and applies directly to main.

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 that needs the database queues behind them.

session_heartbeat is the canary — a trivial UPDATE on one row of sessions, which should never be slow. Measured with skyrl/benchmarks/bench_future_waiting.py (no GPU, no HTTP server) at 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

The poll rate not keeping up with demand is the tell: the pool is the ceiling, so pollers back up and everything sharing the pool waits behind them.

This is also visible in the field. The stuck Tinker Fully Async E2E job logs Session heartbeat failed for 120–220 seconds from its first sampling wave onward, on a run configured for groups_per_batch=512 × group_size=4 = 2048 concurrent samples. I could not reproduce the full 120 s+ on an uncontended box (I got to ~4 s), so I would not claim this as that job's sole cause — the box also has an FSDP trainer, two vLLM engines and a router competing for CPU, and generation was already slow before anything wedged. But the mechanism is the same one, and it is removed here.

How it works

FutureWaiter keeps a dict[request_id, set[asyncio.Future]]. wait() registers a future and awaits it; a single background task polls every 50 ms with one WHERE request_id IN (...) covering every waiter, then resolves them.

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

Details worth reviewing:

• The poller idles properly. With no waiters it blocks on an asyncio.Event rather than spinning. Registration adds to the dict and sets the event with no await in between, so a wakeup cannot be lost — that ordering is load-bearing and commented as such.
• 404 still works. Rows are never deleted, so an id with no row never existed; those waiters get a KeyError, which the endpoint turns into a 404. Previously this was detected on the first poll; now it takes up to one tick (50 ms).
• The poller survives errors. A failed iteration is logged and the loop continues, so a transient database problem cannot permanently wedge every waiter. Callers still have their own timeouts.
• Timeout behaviour is unchanged — still 408 after 300 s, now a named constant (RETRIEVE_FUTURE_TIMEOUT_SECONDS) rather than a literal.
• Chunked at 500 ids per statement, under SQLite's bound-parameter cap.

Risk

This is the highest-risk 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, same body — and test_api.py's integration tests exercise it end to end against a real server subprocess, which is the main evidence here.

Not included

FutureWaiter.notify() — letting the sample-forwarding path hand its result straight to waiters instead of waiting for the poller to rediscover its own write — is deliberately left out, because it needs the result-write changes that are in a separate PR. Without it, forwarded samples are simply discovered on the next 50 ms tick, which is still far better than before.

Testing

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

New in tests/tinker/test_futures.py:

• resolution once a result lands; failed status surfaced distinctly
• KeyError for an unknown request (the 404 path)
• None on timeout (the 408 path)
• multiple waiters on the same request 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 someone reintroduced per-caller polling

Unrelated flake, worth knowing: running the suite twice in a row fails test_api.py::test_training_workflow and ::test_delete_checkpoint with 404s, because the default database_url is a persistent file in the source tree (skyrl/tinker/tinker.db) and consecutive runs contaminate each other. Pre-existing on main. rm -f skyrl/tinker/tinker.db* before rerunning.

🤖 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.

[gemini-code-assist]

In FutureWaiter.stop(), the background poller task is cancelled and awaited. However, any active wait() calls currently awaiting results will be left hanging until their individual timeouts (up to 300 seconds) expire, because the background poller is no longer running to resolve them. This can cause the API server's graceful shutdown (e.g., during redeployments or scaling down) to hang or delay significantly.

To prevent this, we should explicitly cancel all pending waiters in stop() so they fail fast and allow the server to shut down immediately.

    async def stop(self) -> None:
        """Stop the background poller."""
        if self._task is not None:
            self._task.cancel()
            with suppress(asyncio.CancelledError):
                await self._task
            self._task = None
        for waiters in list(self._waiters.values()):
            for waiter in waiters:
                if not waiter.done():
                    waiter.cancel()

[gemini-code-assist]

If the database becomes permanently unavailable or experiences a transient outage, the background poller loop will fail on _poll_once() and immediately log an exception, then sleep for only self._poll_interval_sec (50ms) before retrying. This can rapidly flood the application logs with thousands of stack traces and put unnecessary connection pressure on a struggling database.

Implementing a simple exponential backoff on consecutive failures will mitigate log flooding and allow the database room to recover.

    async def _run(self) -> None:
        consecutive_failures = 0
        while True:
            if not self._waiters:
                # Safe against lost wakeups: registration adds to _waiters and
                # sets the event with no await in between, so an empty dict here
                # means nothing has been registered yet.
                self._wakeup.clear()
                await self._wakeup.wait()
                consecutive_failures = 0
                continue
            try:
                await self._poll_once()
                consecutive_failures = 0
            except asyncio.CancelledError:
                raise
            except Exception:
                consecutive_failures += 1
                sleep_sec = min(self._poll_interval_sec * (2 ** consecutive_failures), 5.0)
                logger.exception("Future poller iteration failed, retrying in %.2fs", sleep_sec)
                await asyncio.sleep(sleep_sec)
                continue
            await asyncio.sleep(self._poll_interval_sec)

[gemini-code-assist]

If a client sends a non-integer or malformed request_id in the RetrieveFutureRequest, calling int(request.request_id) will raise a ValueError. This unhandled exception will propagate up and result in a 500 Internal Server Error.

We should catch ValueError and raise a 400 Bad Request with a clear error message to improve API robustness.

Suggested change

	try:
	result = await req.app.state.future_waiter.wait(int(request.request_id), RETRIEVE_FUTURE_TIMEOUT_SECONDS)
	except KeyError:
	raise HTTPException(status_code=404, detail="Future not found")
	try:
	request_id = int(request.request_id)
	except ValueError:
	raise HTTPException(status_code=400, detail="Invalid request_id format")
	try:
	result = await req.app.state.future_waiter.wait(request_id, RETRIEVE_FUTURE_TIMEOUT_SECONDS)
	except KeyError:
	raise HTTPException(status_code=404, detail="Future not found")