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

skyrl/tinker/api.py

@@ -315,15 +315,19 @@ async def lifespan(app: FastAPI):
     # SkyRL-Train default is colocate_all=True; only opt into forwarding
     # when the operator explicitly sets it to False.
     is_colocated = bool(backend_cfg.get("trainer.placement.colocate_all", True))
+    store_ttls = dict(
+        retrieved_ttl_sec=app.state.engine_config.external_future_retrieved_ttl_sec,
+        completed_ttl_sec=app.state.engine_config.external_future_completed_ttl_sec,
+    )
     if app.state.engine_config.external_inference_url:
-        app.state.external_future_store = ExternalFutureStore()
+        app.state.external_future_store = ExternalFutureStore(**store_ttls)
         await app.state.external_future_store.start()
         app.state.external_inference_client = ExternalInferenceClient(
             app.state.engine_config, app.state.db_engine, app.state.external_future_store
         )
         logger.info(f"External engine configured: {app.state.engine_config.external_inference_url}")
     elif backend_name in ("megatron", "fsdp") and not is_colocated:
-        app.state.external_future_store = ExternalFutureStore()
+        app.state.external_future_store = ExternalFutureStore(**store_ttls)
         await app.state.external_future_store.start()
         app.state.external_inference_client = SkyRLTrainInferenceForwardingClient(
             app.state.engine_config, app.state.db_engine, app.state.external_future_store
@@ -1512,14 +1516,23 @@ async def retrieve_future(request: RetrieveFutureRequest, req: Request):
             types.RequestType(request_type) in PROTO_SERIALIZABLE_REQUEST_TYPES
             and PROTO_CONTENT_TYPE in req.headers.get("accept", "").lower()
         ):
-            async with req.app.state.proto_serialization_lock:
-                content = await asyncio.to_thread(
-                    _serialize_proto_result,
-                    types.RequestType(request_type),
-                    result_data,
-                )
+            # Forwarded samples are stored as proto already and go out as-is;
+            # anything stored as JSON is encoded once here and cached.
+            content = external_future_store.proto_result(request_id) if found_in_memory else None
+            if content is None:
+                async with req.app.state.proto_serialization_lock:
+                    content = await asyncio.to_thread(
+                        _serialize_proto_result,
+                        types.RequestType(request_type),
+                        result_data,
+                    )
+                if found_in_memory:
+                    external_future_store.cache_proto(request_id, content)
             response: Response = Response(content=content, media_type=PROTO_CONTENT_TYPE)
         else:
+            if result_data is None and found_in_memory:
+                # Stored as proto only; a pre-proto client wants JSON.
+                result_data = external_future_store.json_result(request_id)
             response = raw_json_response(result_data)
         # Start the retry-grace clock now that the response is built and about to
         # be sent, so a large result is never evicted mid-delivery -- but only if

skyrl/tinker/config.py

@@ -73,6 +73,21 @@ class EngineConfig(BaseModel):
             "env_var": "SKYRL_FORWARDING_INFERENCE_TIMEOUT_SEC",
         },
     )
+    external_future_retrieved_ttl_sec: float = Field(
+        default=300.0,
+        gt=0,
+        description=(
+            "How long a forwarded sample result stays in memory after it was delivered, so an "
+            "SDK retry after a lost HTTP response still finds it. Must outlast the SDK's worst-case "
+            "re-poll gap (45s poll timeout + up to 30s backoff, twice). Memory for long-output "
+            "rollouts is roughly completion rate x result size x this window."
+        ),
+    )
+    external_future_completed_ttl_sec: float = Field(
+        default=600.0,
+        gt=0,
+        description="How long a completed but never-delivered forwarded sample result stays in memory.",
+    )
     session_cleanup_interval_sec: int = Field(
         default=60,
         description="How often to check for stale sessions (seconds). Set to -1 to disable cleanup.",

skyrl/tinker/external_future_store.py

@@ -9,16 +9,29 @@
 
 from skyrl.tinker import types
 from skyrl.tinker.db_models import RequestStatus
+from skyrl.tinker.proto_serialization import sample_output_json_from_proto
 from skyrl.utils.log import logger
 
 
+@dataclass
+class PreparedResult:
+    """A completed sample result already in wire form.
+
+    Forwarded samples are encoded to proto once, straight from the decoded vLLM
+    body; JSON text is produced only if a pre-proto client asks for it.
+    """
+
+    proto: bytes | None = None
+    json: str | None = None
+
+
 @dataclass
 class ExternalFuture:
     request_id: int
     model_id: str | None
-    request_data: dict
     status: RequestStatus = RequestStatus.PENDING
     result_data: str | None = None
+    result_proto: bytes | None = None
     created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
     completed_at: datetime | None = None
     retrieved_at: datetime | None = None
@@ -54,7 +67,11 @@ class ExternalFutureStore:
     # Pending entries whose forwarding task died without completing them.
     _PENDING_TTL_SECONDS = 3600.0
 
-    def __init__(self):
+    def __init__(self, *, retrieved_ttl_sec: float | None = None, completed_ttl_sec: float | None = None):
+        if retrieved_ttl_sec is not None:
+            self._RETRIEVED_TTL_SECONDS = retrieved_ttl_sec
+        if completed_ttl_sec is not None:
+            self._COMPLETED_TTL_SECONDS = completed_ttl_sec
         self._entries: dict[int, ExternalFuture] = {}
         # Boot-epoch id space: each server process starts below every id an
         # earlier process could plausibly have handed out (2^20 ids per
@@ -70,13 +87,29 @@ async def start(self) -> None:
     def create(self, model_id: str | None, request_data: BaseModel) -> int:
         request_id = self._next_request_id
         self._next_request_id -= 1
-        self._entries[request_id] = ExternalFuture(
-            request_id=request_id,
-            model_id=model_id,
-            request_data=request_data.model_dump(mode="json"),
-        )
+        self._entries[request_id] = ExternalFuture(request_id=request_id, model_id=model_id)
         return request_id
 
+    def proto_result(self, request_id: int) -> bytes | None:
+        """Proto wire bytes for a completed result, if it has them."""
+        entry = self._entries.get(request_id)
+        return entry.result_proto if entry is not None else None
+
+    def cache_proto(self, request_id: int, proto: bytes) -> None:
+        """Keep a proto encoding produced at retrieval so retries skip re-encoding."""
+        entry = self._entries.get(request_id)
+        if entry is not None:
+            entry.result_proto = proto
+
+    def json_result(self, request_id: int) -> str | None:
+        """JSON text for a completed result, deriving it from proto on first use."""
+        entry = self._entries.get(request_id)
+        if entry is None:
+            return None
+        if entry.result_data is None and entry.result_proto is not None:
+            entry.result_data = sample_output_json_from_proto(entry.result_proto)
+        return entry.result_data
+
     async def wait(self, request_id: int, timeout: float) -> tuple[RequestStatus, types.RequestType, str | None] | None:
         entry = self._entries.get(request_id)
         if entry is None:
@@ -99,13 +132,17 @@ def mark_retrieved(self, request_id: int) -> None:
         if entry is not None:
             entry.retrieved_at = datetime.now(timezone.utc)
 
-    async def complete(self, request_id: int, result_data: BaseModel, status: RequestStatus) -> None:
+    async def complete(self, request_id: int, result_data: BaseModel | PreparedResult, status: RequestStatus) -> None:
         entry = self._entries.get(request_id)
         if entry is None:
             # Swept as abandoned before the forwarding task finished.
             logger.warning("External future %s was evicted before its result arrived — dropping", request_id)
             return
-        entry.result_data = result_data.model_dump_json()
+        if isinstance(result_data, PreparedResult):
+            entry.result_data = result_data.json
+            entry.result_proto = result_data.proto
+        else:
+            entry.result_data = result_data.model_dump_json()
         entry.status = status
         entry.completed_at = datetime.now(timezone.utc)
         entry.event.set()

skyrl/tinker/extra/skyrl_train_inference_forwarding.py

@@ -16,7 +16,11 @@
 from skyrl.tinker import types
 from skyrl.tinker.config import EngineConfig
 from skyrl.tinker.db_models import EngineStateDB, FutureDB, RequestStatus
-from skyrl.tinker.external_future_store import ExternalFutureStore
+from skyrl.tinker.external_future_store import ExternalFutureStore, PreparedResult
+from skyrl.tinker.proto_serialization import (
+    sample_output_json_from_proto,
+    serialize_sample_output,
+)
 from skyrl.utils.log import logger
 
 
@@ -142,12 +146,15 @@ async def call_and_store_result(
                 logger.warning("FutureDB row %s missing on completion write — skipping", request_id)
                 return
             # `result_data` is a text column holding pre-serialized JSON.
-            future.result_data = result.model_dump_json()
+            if isinstance(result, PreparedResult):
+                future.result_data = result.json or sample_output_json_from_proto(result.proto)
+            else:
+                future.result_data = result.model_dump_json()
             future.status = status
             future.completed_at = datetime.now(timezone.utc)
             await session.commit()
 
-    async def _forward_with_retry(self, sample_req, model_id: str, *, base_model: str | None) -> types.SampleOutput:
+    async def _forward_with_retry(self, sample_req, model_id: str, *, base_model: str | None) -> PreparedResult:
         # Retry only failures where the request demonstrably did not execute:
         # connect-phase errors and 5xx rejections from the router. Read and
         # write failures are ambiguous: vLLM may still be executing the
@@ -179,9 +186,7 @@ async def _forward_with_retry(self, sample_req, model_id: str, *, base_model: st
                 "the SKYRL_FORWARDING_INFERENCE_TIMEOUT_SEC environment variable."
             ) from e
 
-    async def _forward(
-        self, proxy_url: str, sample_req, model_id: str, *, base_model: str | None
-    ) -> types.SampleOutput:
+    async def _forward(self, proxy_url: str, sample_req, model_id: str, *, base_model: str | None) -> PreparedResult:
         # model_id matches the LoRA name registered with vLLM during
         # save_weights_for_sampler; base_model is used for non-LoRA sampling.
         model_name = base_model if base_model else model_id
@@ -266,16 +271,8 @@ async def _forward(
             # Tinker's stop_reason is Literal["stop", "length"]; vLLM emits a wider set.
             finish_reason = choice.get("finish_reason")
             stop_reason = "stop" if finish_reason in ("stop", "stop_token") else "length"
-            sequences.append(
-                types.GeneratedSequence(
-                    tokens=tokens,
-                    logprobs=logprobs,
-                    stop_reason=stop_reason,
-                )
-            )
+            sequences.append((stop_reason, tokens, logprobs))
 
-        return types.SampleOutput(
-            sequences=sequences,
-            prompt_logprobs=prompt_logprobs,
-            topk_prompt_logprobs=topk,
-        )
+        # Encode straight to the proto wire form the SDK retrieves; no pydantic
+        # model or JSON text is built for the result (see PreparedResult).
+        return PreparedResult(proto=serialize_sample_output(sequences, prompt_logprobs, topk))

skyrl/tinker/proto_serialization.py

@@ -19,8 +19,10 @@
 """
 
 import base64
+from collections.abc import Iterable, Sequence
 
 import numpy as np
+import orjson
 from tinker.proto import tinker_public_pb2 as pb
 
 from skyrl.tinker import types
@@ -124,24 +126,44 @@ def serialize_result(request_type: types.RequestType, result_data: dict) -> byte
 
 def _serialize_sample_output(result_data: dict) -> bytes:
     output = types.SampleOutput.model_validate(result_data)
+    return serialize_sample_output(
+        [(seq.stop_reason, seq.tokens, seq.logprobs) for seq in output.sequences],
+        output.prompt_logprobs,
+        output.topk_prompt_logprobs,
+    )
+
+
+def serialize_sample_output(
+    sequences: Iterable[tuple[str, Sequence[int], Sequence[float]]],
+    prompt_logprobs: Sequence[float | None] | None,
+    topk_prompt_logprobs: Sequence[Sequence[tuple[int, float]] | None] | None,
+) -> bytes:
+    """Build ``SampleResponse`` wire bytes from plain Python data.
+
+    ``sequences`` holds ``(stop_reason, tokens, logprobs)`` per sequence. This is
+    the hot path for forwarded samples: the vLLM body is decoded once and
+    encoded straight to proto, with no pydantic model and no JSON text in
+    between (each of which costs about as much as this whole function for a
+    32k-token result).
+    """
     proto = pb.SampleResponse()
 
-    for seq in output.sequences:
+    for stop_reason, tokens, logprobs in sequences:
         proto.sequences.append(
             pb.SampledSequence(
-                stop_reason=_STOP_REASON_TO_PROTO[seq.stop_reason],
-                tokens=np.asarray(seq.tokens, dtype=np.int32).tobytes(),
-                logprobs=np.asarray(seq.logprobs, dtype=np.float32).tobytes(),
+                stop_reason=_STOP_REASON_TO_PROTO[stop_reason],
+                tokens=np.asarray(tokens, dtype=np.int32).tobytes(),
+                logprobs=np.asarray(logprobs, dtype=np.float32).tobytes(),
             )
         )
 
-    if output.prompt_logprobs is not None:
+    if prompt_logprobs is not None:
         proto.prompt_logprobs = np.array(
-            [np.nan if lp is None else lp for lp in output.prompt_logprobs], dtype=np.float32
+            [np.nan if lp is None else lp for lp in prompt_logprobs], dtype=np.float32
         ).tobytes()
 
-    if output.topk_prompt_logprobs is not None:
-        rows = output.topk_prompt_logprobs
+    if topk_prompt_logprobs is not None:
+        rows = topk_prompt_logprobs
         # k is not recorded in the result, so recover it from the widest row.
         # With every row undefined, use k=1 so prompt_length stays encoded
         # (the client maps fully-masked rows back to None).
@@ -164,6 +186,49 @@ def _serialize_sample_output(result_data: dict) -> bytes:
     return proto.SerializeToString()
 
 
+_PROTO_TO_STOP_REASON = {value: key for key, value in _STOP_REASON_TO_PROTO.items()}
+
+
+def sample_output_json_from_proto(proto_bytes: bytes) -> str:
+    """Inverse of :func:`serialize_sample_output`, as ``SampleOutput`` JSON text.
+
+    Serves clients that predate proto results (SDK < 0.25) from a result that
+    was stored as proto. Logprobs come back at float32 precision, the same
+    values proto clients receive.
+    """
+    proto = pb.SampleResponse.FromString(proto_bytes)
+    sequences = [
+        {
+            "stop_reason": _PROTO_TO_STOP_REASON[seq.stop_reason],
+            "tokens": np.frombuffer(seq.tokens, dtype=np.int32).tolist(),
+            "logprobs": np.frombuffer(seq.logprobs, dtype=np.float32).tolist(),
+        }
+        for seq in proto.sequences
+    ]
+
+    prompt_logprobs = None
+    if proto.prompt_logprobs:
+        values = np.frombuffer(proto.prompt_logprobs, dtype=np.float32)
+        prompt_logprobs = [None if np.isnan(value) else float(value) for value in values]
+
+    topk = None
+    if proto.HasField("topk_prompt_logprobs"):
+        block = proto.topk_prompt_logprobs
+        shape = (block.prompt_length, block.k)
+        token_ids = np.frombuffer(block.token_ids, dtype=np.int32).reshape(shape)
+        logprobs = np.frombuffer(block.logprobs, dtype=np.float32).reshape(shape)
+        masked = (token_ids == _TOPK_MASK_TOKEN_ID) & (logprobs == np.float32(_TOPK_MASK_LOGPROB))
+        topk = []
+        for row_ids, row_lps, row_mask in zip(token_ids, logprobs, masked):
+            row = [[int(t), float(lp)] for t, lp, m in zip(row_ids, row_lps, row_mask) if not m]
+            # A fully masked row encodes an undefined position (None).
+            topk.append(row or None)
+
+    return orjson.dumps(
+        {"sequences": sequences, "prompt_logprobs": prompt_logprobs, "topk_prompt_logprobs": topk}
+    ).decode()
+
+
 def _serialize_forward_backward_output(result_data: dict) -> bytes:
     output = types.ForwardBackwardOutput.model_validate(result_data)
     proto = pb.ForwardBackwardOutput()

tests/tinker/test_external_future_store.py

@@ -277,6 +277,12 @@ async def wait(self, request_id, timeout):
         def mark_retrieved(self, request_id):
             pass
 
+        def proto_result(self, request_id):
+            return None
+
+        def cache_proto(self, request_id, proto):
+            pass
+
     def serialize_result_in_thread(request_type, result_data):
         nonlocal active_serializations, max_active_serializations
         serialization_thread_ids.append(threading.get_ident())