Skip to content

fix(tinker): avoid SQLite pool exhaustion under rollout load - #2097

Open
hershg wants to merge 6 commits into
NovaSky-AI:mainfrom
hershg:hersh/external-future-store-upstream
Open

fix(tinker): avoid SQLite pool exhaustion under rollout load#2097
hershg wants to merge 6 commits into
NovaSky-AI:mainfrom
hershg:hersh/external-future-store-upstream

Conversation

@hershg

@hershg hershg commented Aug 27, 2026

Copy link
Copy Markdown

Summary

  • keep externally forwarded sample futures off SQLite's request hot path while preserving durable terminal records
  • batch terminal persistence and serialize only the SQLite writes that remain
  • cache immutable sampler validation, own forwarding tasks through shutdown, and preserve Tinker 0.25 protobuf responses

Why

Tau creates 512 concurrent samples alongside training futures and heartbeats. Stock SkyRL repeatedly validates the same sampler checkpoint and creates/completes each future through a SQLite pool of 5+10 connections. Under the real mix, the pool exhausts and accepted training stalls.

This is a source-level runtime fix. It does not enlarge the pool or add model-, accelerator-, or workload-specific behavior.

Comparison

All arms were pinned to their exact heads and exercised with the same four-wave gate: 512 model-path samples, 512 forward/backward futures, and 32 heartbeats per wave.

Arm Exact head Result
Stock consumer source 27ae52bf fails in 4.01s with QueuePool size 5 overflow 10
#2059 92c3628d fails in wave 1; sample validation, training futures, and heartbeats still contend on SQLite
#2061 2760b587 samples 2,048/2,048, but training 511/2,048 and heartbeats 7/128; queue rejects sequence 512 and strands later work
This PR 20db38cf samples 2,048/2,048, training 2,048/2,048, heartbeats 128/128; repeated retrieval and terminal persistence pass

#2059's negative in-memory IDs are useful, but its terminal results are memory-only, its external path returns JSON rather than protobuf, and its forwarding tasks are not drained at shutdown. #2061 removes more SQLite traffic, but the current queue/Unix-socket design adds substantially more code and fails admission, reconnect, cache-invalidation, and shutdown contracts. This PR is the smallest arm that passes the representative workload and preserves durable semantics.

Lifecycle guarantees

  • completion becomes visible only after terminal persistence succeeds
  • persistence failures propagate to waiters
  • persisted entries are released from memory and remain retrievable from SQLite
  • forwarding tasks finish before the inference client and future store close
  • sampler deletion synchronizes with validation and invalidates the cache
  • write serialization applies to SQLite, not Postgres
  • the background engine still stops if terminal persistence fails during shutdown

Evidence

  • exact sustained gate: 1 passed in 12.23s
  • focused affected suite: 18 passed, 4 skipped in 18.10s
  • adversarial lifecycle/durability selection: 7 passed in 5.17s
  • dedicated external-future suite: 9 passed in 9.50s
  • full non-integration Tinker suite: 99 passed, 23 skipped in 390.37s
  • pre-commit: Ruff, Black, and gitleaks passed
  • git diff --check 59d4daed...HEAD passed

Exact live consumer proof: deployment XID 1048275 cold-deployed image sha256:efb03cc344ea373f514a78859f4c0b9a91d3f58b187ab2c201ed8699e499ce11 from SkyRL source 27d94205. That source contains this PR's first four commits unchanged and the final lifecycle commit rebased as 747b5bd1; both lifecycle commits have stable patch ID 9668ff4d.

Training XID 1048276 then completed the full real workload on that exact image:

  • ten optimizer steps, including repeated 512-concurrent rollout waves, packed forward/backward, and weight publication
  • durable checkpoints at steps 0, 5, and 10
  • ten 50-example evals, 500/500 accounted with zero stale, timed-out, errored, or missing trajectories
  • no QueuePool/SQLite failure, missing future, lost heartbeat, traceback, or service restart

The same mechanism previously completed ten optimizer steps in XIDs 1048178 and 1048230, using precursor source abc014a5. XID 1048276 closes the exact-current-patch integration gap. Current head 20db38cf additionally regression-tests that a terminal-persistence failure is still propagated while the background engine is always stopped.

Tests

uv run --isolated --frozen --extra dev --extra tinker --extra ray --extra jax \
  pytest -q tests/tinker/test_external_future_store.py

uv run --no-sync pytest -q tests/tinker -m 'not integration'

pre-commit run --all-files --config .pre-commit-config.yaml

Note

Medium Risk
Changes async request lifecycle, SQLite write ordering, and shutdown for the hot rollout path; failures surface as persistence errors to waiters rather than silent loss of durable futures.

Overview
Under high concurrent rollout load (many samples plus training futures and heartbeats), the Tinker API was exhausting the small SQLite connection pool because every forwarded sample created and completed futures and repeatedly validated the same sampler checkpoint on the DB hot path.

This PR adds ExternalFutureStore for non-colocated SkyRL-Train inference forwarding: samples get in-memory futures with negative IDs, retrieve_future waits on the store first, and terminal rows are batched into SQLite only after inference completes—waiters are not released until persistence succeeds. SkyRLTrainInferenceForwardingClient completes through the store instead of per-request FutureDB updates.

Remaining SQLite writes on SQLite are serialized via db_write_lock (heartbeats, forward_backward creation); Postgres keeps parallel writes. asample caches sampling-session model resolution and validates sampler checkpoints once per (model_id, checkpoint_id); deletion shares the validation lock and clears that cache.

Shutdown now drains tracked forwarding tasks before closing the inference client and future store, then stops the background engine even if external close fails.

Reviewed by Cursor Bugbot for commit 20db38c. Bugbot is set up for automated code reviews on this repo. Configure here.

@hershg
hershg marked this pull request as ready for review August 27, 2026 23:18

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces an ExternalFutureStore to keep forwarded sample futures off the database hot path by managing them in memory and persisting them asynchronously in batches. It also implements caching for sampling models and sampler checkpoint validations, and introduces database write serialization locks (particularly for SQLite) to reduce database contention. I have no feedback to provide as there are no review comments to assess.

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

Fix All in Cursor

Reviewed by Cursor Bugbot for commit 2a507ce. Configure here.

Comment thread skyrl/tinker/api.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant