[perf][tinker] Don't deserialize request payloads the scheduler discards

[avigyabb]

What

The engine's scheduling queries selected whole FutureDB rows, so every pending request's request_data was deserialized on every poll — including requests parked behind an optim_step/load_weights barrier, which are then filtered out and thrown away.

A forward_backward with 4×512 tokens is ~76 KiB of JSON. A backlog of a few hundred parked requests therefore meant tens of MB of wasted decode per 100 ms tick, which keeps the engine CPU-bound and slow to dispatch. That backlog is the multi-LoRA steady state: some adapters stepping while others have passes stacked up behind a barrier.

The change

The three finders now select ids only, filter exactly as before, and fetch request_data for just the requests they return. Sample batching needs one field from the payload, so checkpoint_id is extracted in the database rather than loading whole prompts to reach it.

Per finder the change is two lines — the column list, and a payload fetch after the filter:

-        select(FutureDB)
+        select(FutureDB.request_id, FutureDB.model_id)
...
+        payloads = self._load_request_payloads(session, [request_id for request_id, _ in batchable])

Plus one small helper, _load_request_payloads, so that fetch isn't written three times.

Scheduling semantics are unchanged: same barrier logic, same request_id ordering, same return shapes. The existing barrier regression tests pass untouched.

Where the cost was

Measured on 512 pending rows of ~76 KiB, breaking down what the old full-row query was paying for:

Python JSON deserialization	499 ms (83%)
ORM object construction	73 ms (12%)
reading payload text off disk	30 ms (5%)
the id-only query itself	1.6 ms (0.3%)

Note how little of it is I/O — only 30 ms of 602 ms is SQLite reading text. The rest is Python turning that text into dicts and lists of numbers. That's why this showed up as engine CPU rather than disk, and why it starves other processes on a shared box.

A scan that can dispatch nothing drops from ~285 ms to ~2 ms. A scan where everything is dispatchable is roughly unchanged, because those payloads are genuinely needed. The saving is proportional to how much of the queue is blocked — which for multi-LoRA is most of it.

Deliberately kept small

An earlier version of this (#1976) also replaced the per-finder queries with one shared metadata scan, moved barrier computation from SQL into Python, and added a covering index. Measuring the pieces separately showed that was poor value: the shared scan is worth about 1 ms on top of this (0.4% of the win), and the covering index is below noise once payloads are skipped. All of it is left out here, and _find_destructive_barriers keeps its existing GROUP BY query untouched.

I'd close #1976 and #1977 in favour of this if it looks right.

One behaviour note

json_extract returns NULL for a missing key, where the previous op.request_data["checkpoint_id"] would have raised KeyError. checkpoint_id is required on SampleInput and always present in model_dump(), so this is unreachable today; a missing value is now treated like the base-model case (always batch-compatible) rather than crashing the scan.

Testing

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

Three tests added to tests/tinker/test_engine.py, covering what previously had none:

• find_batchable_model_passes holds back passes behind a barrier while another model's pass still dispatches, and the returned payload still parses
• find_batchable_sample keeps one checkpoint_id per model, exercising the json_extract path
• a backlog larger than the per-statement id chunk still resolves fully

🤖 Generated with Claude Code