Skip to content

[Tinker API] Keep forwarded sample futures in memory - #2059

Draft
yapdianang wants to merge 5 commits into
NovaSky-AI:mainfrom
yapdianang:dydx/tinker-in-memory-futures-current
Draft

[Tinker API] Keep forwarded sample futures in memory#2059
yapdianang wants to merge 5 commits into
NovaSky-AI:mainfrom
yapdianang:dydx/tinker-in-memory-futures-current

Conversation

@yapdianang

Copy link
Copy Markdown
Contributor

Summary

  • Keep forwarded sample futures in process memory.
  • Return JSON and protobuf sample results through the same future API.
  • Remove a completed result in constant time when the client retrieves it.
  • Keep training futures in the database.

Depends on #2043.

Before / After

Path Before After
Forwarded sample submit Create an external future row in SQLite Create an in-memory future ID
Forwarded sample completion Write the sample result to SQLite Complete the in-memory future
Forwarded sample retrieval Read and then clean up the database row Pop the result from memory in constant time
Training operations Database-backed Database-backed

This removes sample-result database traffic from the high-concurrency rollout path. The store has a fixed capacity and rejects new futures when full.

Testing

uv run --extra dev --extra jax --extra tinker pytest \
  tests/tinker/test_future_waiting.py \
  tests/tinker/test_in_memory_future_store.py \
  tests/tinker/test_proto_serialization.py \
  tests/tinker/skyrl_train/test_async_sample_routing.py \
  tests/tinker/test_api.py
uv run --extra dev pre-commit run --all-files

The focused API tests pass: 45 passed and 4 skipped. All pre-commit hooks pass.

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