[tinker] Hold 131k concurrent samples incl. long outputs: aiohttp forwarding, proto-once results, cheap access log, undelivered-result fix, load harness

skyrl/tinker/api.py

@@ -308,15 +308,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
@@ -929,12 +933,17 @@ class WeightsInfoResponse(BaseModel):
 
 class ClientConfigResponse(BaseModel):
     pjwt_auth_enabled: bool = False
+    # Cap on in-flight samples per SamplingClient; the SDK applies it as its
+    # sampling RetryConfig.max_connections.
+    sample_max_concurrent_requests: int = 2000
 
 
 @app.post("/api/v1/client/config", response_model=ClientConfigResponse)
-async def client_config():
-    """Stub for tinker SDK client_config handshake."""
-    return ClientConfigResponse()
+async def client_config(req: Request):
+    """Tinker SDK client_config handshake: server-side flags for the client."""
+    return ClientConfigResponse(
+        sample_max_concurrent_requests=req.app.state.engine_config.sample_max_concurrent_requests,
+    )
 
 
 @app.get("/api/v1/healthz", response_model=HealthResponse)
@@ -1505,18 +1514,32 @@ 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.
-        if found_in_memory:
+        # be sent, so a large result is never evicted mid-delivery -- but only if
+        # this client is still there to receive it. The SDK abandons a poll after
+        # 45s and retries the same request_id; if the result lands after that,
+        # this handler wakes on a dead connection (uvicorn drops the send
+        # silently) and starting the short clock here would let the sweeper
+        # evict a result nobody received, turning the retry into a 404.
+        if found_in_memory and not await req.is_disconnected():
             external_future_store.mark_retrieved(request_id)
         return response
 

skyrl/tinker/config.py

@@ -59,18 +59,46 @@ class EngineConfig(BaseModel):
         json_schema_extra={"argparse_type": lambda v: None if v == "None" else int(v)},
     )
     forwarding_inference_timeout_sec: float = Field(
-        default=300.0,
+        default=2048.0,
         gt=0,
         description=(
             "Read timeout in seconds for API-side requests forwarded to the "
             "SkyRL-Train-managed inference engine. This must cover time spent "
-            "queued behind other requests as well as generation time."
+            "queued behind other requests as well as generation time: with the "
+            "default unlimited connection count a large rollout burst waits inside "
+            "vLLM's queue, and 128x128 bursts routinely exceed 300s there."
         ),
         json_schema_extra={
             "argparse_type": float,
             "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.",
+    )
+    sample_max_concurrent_requests: int = Field(
+        default=2000,
+        gt=0,
+        description=(
+            "Advertised to the Tinker SDK via /api/v1/client/config as the maximum "
+            "number of sample requests one SamplingClient keeps in flight; the SDK "
+            "queues the rest client-side. Each in-flight sample costs the API server "
+            "one open long-poll connection and one retrieve_future re-poll every 45s. "
+            "2000 is the SDK's own default."
+        ),
+    )
     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,31 @@
 
 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
+    # At most one of these is populated at completion; the other is derived
+    # lazily on the first request for it and then cached for retries.
     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
@@ -44,15 +59,23 @@ class ExternalFutureStore:
     # retry following a lost HTTP response still finds it. Measured from
     # delivery (mark_retrieved), never from the in-store read: a large result
     # can spend minutes being serialized and sent, and starting the clock at
-    # read would evict it mid-delivery.
-    _RETRIEVED_TTL_SECONDS = 120.0
+    # read would evict it mid-delivery. The SDK re-polls after a 45s client
+    # timeout plus up to 30s of backoff, so two consecutive misses span 150s;
+    # the grace has to outlast that.
+    _RETRIEVED_TTL_SECONDS = 300.0
     # Completed but not yet delivered — governs the read/serialize/send window
     # and clients that never come back.
     _COMPLETED_TTL_SECONDS = 600.0
     # 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):
+        # Retention after delivery is the dominant memory term for long-output
+        # rollouts (results x retrieved TTL), so operators can tune it.
+        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
@@ -66,15 +89,33 @@ async def start(self) -> None:
         self._sweeper = asyncio.create_task(self._sweep_loop())
 
     def create(self, model_id: str | None, request_data: BaseModel) -> int:
+        # The request body is not retained: nothing reads it back on this
+        # path, and a long prompt would cost ~100KB per pending entry.
         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:
@@ -97,13 +138,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()