[tinker] Encode forwarded sample results to proto once and serve them as-is

[avigyabb]

Stack 5/7. Long-output rollouts made per-result payload work the API server's main cost.

Before. For a 32k-token result (356 KB JSON) the forwarding path spent 4 ms (orjson decode, pydantic validate, pydantic JSON dump) and the proto path the SDK >= 0.25 uses on retrieve_future spent another 7 ms (stdlib json.loads + proto build) inside a single global lock, capping proto delivery near 150 results/s regardless of concurrency. The proto build holds the GIL, so the thread hop bought nothing.

Change.

• The forwarding client decodes the vLLM body once and encodes straight to SampleResponse wire bytes (serialize_sample_output, shared with the validated path and pinned to it byte-for-byte by tests). No pydantic model or JSON text is built for the result.
• ExternalFutureStore keeps the proto bytes (8 bytes/token, 26% smaller than the JSON text); retrieve_future passes them through for proto clients and derives JSON lazily, cached, for pre-0.25 clients (sample_output_json_from_proto, float32 logprobs). Results stored as JSON (DB path, errors) keep the encode-in-thread path, now cached per entry.
• Pending entries no longer retain the request body (never read back on this path; ~100 KB per entry for long prompts).
• Store TTLs are EngineConfig fields (external_future_retrieved_ttl_sec, external_future_completed_ttl_sec): retention after delivery is the dominant memory term, roughly completion rate x result size x window.

Results (proto client, 2048 cap): 512 concurrent 32k-token results 7.5 s -> 3.5 s server CPU, 73 -> 167 results/s. 131072 requests with 8k-token results, 5 s engine queueing and SDK-style 45 s re-polls: 131072/131072, 0 failures, 267/s, 8.9 GB peak RSS. 32768 x 32k tokens: 32768/32768, 163/s, 9.6 GB.

Stack (each PR retargets to main as the one below merges)

1. [tinker] Route uvicorn's access log to a plain handler instead of Rich #2160
2. [tinker] Forward samples with aiohttp instead of httpx #2161
3. [tinker] Keep an undelivered sample result alive for the SDK's retry #2162
4. [tinker] Survive completion bursts at the socket layer (accept backlog, keep-alive) #2163
5. [tinker] Encode forwarded sample results to proto once and serve them as-is #2164
6. [tinker] Decode vLLM completion bodies straight into numpy with pysimdjson #2165
7. [tinker] Load harness for the API server's sampling path at 131k concurrency #2166

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes the high-throughput sample retrieval and forwarding path; incorrect caching or TTL behavior could cause 404s on SDK retries or extra memory use, but scope is limited to external/forwarded samples with tests covering proto/JSON paths.

Overview
Forwarded inference samples no longer build a Pydantic SampleOutput or JSON on the hot path. The SkyRL-Train forwarding client decodes the vLLM response once and writes SampleResponse proto bytes via shared serialize_sample_output, wrapped in a new PreparedResult.

ExternalFutureStore keeps proto (and optional JSON) per future, drops retained request bodies for pending entries, and exposes proto_result / cache_proto / json_result so retrieve_future can return proto as-is for SDK ≥0.25, lazily derive JSON for older clients, and cache one-time proto encoding for JSON-backed results. Store retention TTLs are configurable on EngineConfig (external_future_retrieved_ttl_sec, external_future_completed_ttl_sec) and passed in at startup.

proto_serialization adds the direct serialize_sample_output helper and sample_output_json_from_proto for backward-compatible JSON. New tests pin byte-for-byte equivalence with the validated path and cover retrieve/cache behavior.

^{Reviewed by Cursor Bugbot for commit 0539b22. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request implements a performance optimization for forwarded sample results by encoding them directly to Protobuf wire format (via PreparedResult) instead of serializing to Pydantic models or JSON text. Pre-proto clients requesting JSON will have the JSON lazily derived from the Protobuf representation and cached. Additionally, configurable TTL settings are introduced for retrieved and completed futures. The review feedback highlights a potential TypeError in skyrl_train_inference_forwarding.py if result.proto is None when attempting to decode it, suggesting a defensive check to prevent this issue.

[greptile-apps]

Confidence Score: 4/5

The PR appears safe to merge, with a non-blocking concurrency issue that can redundantly serialize the same large result during overlapping retrievals.

The new fast path correctly routes stored protobuf and legacy JSON responses, but its cache check occurs outside the serialization lock, allowing concurrent waiters that miss together to repeat expensive serialization sequentially.

Files Needing Attention: skyrl/tinker/api.py

Important Files Changed

Filename	Overview
skyrl/tinker/api.py	Wires configurable store TTLs into startup and adds protobuf pass-through and caching; the pre-lock cache lookup permits duplicate serialization by overlapping waiters.
skyrl/tinker/config.py	Adds validated positive TTL configuration for delivered and undelivered external future results.
skyrl/tinker/external_future_store.py	Introduces prepared protobuf results, lazy JSON conversion, per-entry representation caching, and removes unused request-body retention.
skyrl/tinker/extra/skyrl_train_inference_forwarding.py	Converts decoded vLLM sample output directly into protobuf bytes while retaining JSON persistence for the database path.
skyrl/tinker/proto_serialization.py	Extracts direct SampleResponse serialization and adds an inverse JSON conversion for legacy clients.
tests/tinker/test_external_future_store.py	Updates the test store interface for protobuf cache access.
tests/tinker/test_sample_result_fast_path.py	Covers byte equivalence, conversion, pass-through, sequential caching, and TTL behavior, but not concurrent cache misses.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[vLLM sample response] --> B[Decode response body]
  B --> C[Serialize SampleResponse once]
  C --> D[ExternalFutureStore]
  D -->|Proto client| E[Return stored protobuf bytes]
  D -->|Legacy JSON client| F[Convert protobuf to JSON]
  F --> G[Cache and return JSON]
  H[JSON-backed result] --> D
  D -->|First proto retrieval| I[Serialize under global lock]
  I --> J[Cache and return protobuf]

Loading

_{Reviews (1): Last reviewed commit: "Clean up comments in ExternalFuture clas..." | Re-trigger Greptile}