[tinker] Forward samples with aiohttp instead of httpx

skyrl/tinker/api.py

@@ -32,6 +32,7 @@
 from sqlmodel import SQLModel, func, select
 from sqlmodel.ext.asyncio.session import AsyncSession
 
+from skyrl.env_vars import SKYRL_HTTP_CONNECTION_LIMIT
 from skyrl.tinker import types
 from skyrl.tinker.config import EngineConfig, add_model, config_to_argv
 from skyrl.tinker.db_models import (
@@ -71,6 +72,12 @@
 # How long retrieve_future waits for a result before returning 408
 RETRIEVE_FUTURE_TIMEOUT_SECONDS = 300
 
+# Idle keep-alive for client connections. Under a burst of completions the
+# event loop can be busy for many seconds; with uvicorn's 5s default every
+# idle SDK connection is closed during such a burst and all clients reconnect
+# at once, overflowing the accept backlog. Hold connections across bursts.
+HTTP_KEEP_ALIVE_TIMEOUT_SECONDS = 75
+
 # How often poll_futures looks for newly finished requests. A single query
 # covers every waiter, so this can stay tight without the load scaling up with
 # the number of in-flight requests.
@@ -1515,8 +1522,13 @@ async def retrieve_future(request: RetrieveFutureRequest, req: Request):
         else:
             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
 
@@ -1849,4 +1861,13 @@ async def root():
     # Store config in app.state so lifespan can access it
     app.state.engine_config = engine_config
 
-    uvicorn.run(app, host=args.host, port=args.port, log_config=get_uvicorn_log_config())
+    uvicorn.run(
+        app,
+        host=args.host,
+        port=args.port,
+        log_config=get_uvicorn_log_config(),
+        # Pending connections queue in the kernel while the loop is busy instead
+        # of being refused (effective value is capped by net.core.somaxconn).
+        backlog=SKYRL_HTTP_CONNECTION_LIMIT,
+        timeout_keep_alive=HTTP_KEEP_ALIVE_TIMEOUT_SECONDS,
+    )

skyrl/tinker/config.py

@@ -59,12 +59,14 @@ 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,

skyrl/tinker/external_future_store.py

@@ -44,8 +44,10 @@ 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

skyrl/tinker/extra/skyrl_train_inference_forwarding.py

@@ -7,7 +7,8 @@
 import asyncio
 from datetime import datetime, timezone
 
-import httpx
+import aiohttp
+import orjson
 from sqlmodel.ext.asyncio.session import AsyncSession
 
 from skyrl.backends.renderer import render_model_input
@@ -19,6 +20,13 @@
 from skyrl.utils.log import logger
 
 
+class TransientInferenceError(RuntimeError):
+    """A 5xx from vllm-router/vLLM: the request was rejected, not executed, so it is safe to retry."""
+
+
+_ROUTER_CONNECT_TIMEOUT_SECONDS = 60.0
+
+
 class SkyRLTrainInferenceForwardingClient:
     """Forwards EXTERNAL sample requests to the SkyRL-Train-managed vLLM."""
 
@@ -36,27 +44,48 @@ def __init__(
         self.external_future_store = external_future_store
         self._cached_proxy_url: str | None = None
         self._cache_lock = asyncio.Lock()
-        # Backpressure layered: httpx pool -> vllm-router -> vLLM max_num_seqs.
-        # Default `forwarding_inference_max_connections=None` is unlimited;
-        # the only cost is file descriptors (raise `ulimit -n` accordingly).
-        max_conn = engine_config.forwarding_inference_max_connections
-        max_keepalive = max(max_conn // 4, 32) if max_conn is not None else None
-        self._http_client: httpx.AsyncClient = httpx.AsyncClient(
-            timeout=httpx.Timeout(
-                connect=10.0,
-                read=engine_config.forwarding_inference_timeout_sec,
-                write=300.0,
-                pool=300.0,
-            ),
-            limits=httpx.Limits(
-                max_connections=max_conn,
-                max_keepalive_connections=max_keepalive,
-            ),
-        )
+        # Created on first use so it binds to the serving event loop.
+        self._session: aiohttp.ClientSession | None = None
+
+    def _get_session(self) -> aiohttp.ClientSession:
+        """Return the shared aiohttp session, creating it on first use.
+
+        Backpressure is layered: connector limit -> vllm-router -> vLLM
+        max_num_seqs. Default `forwarding_inference_max_connections=None` is
+        unlimited; the only cost is file descriptors (raise `ulimit -n`
+        accordingly). Requests beyond the limit wait in the connector's FIFO
+        queue with no deadline, so a backlog of many thousands of samples
+        drains at the engine's pace instead of failing.
+
+        aiohttp rather than httpx: httpcore's pool rescans every connection
+        for every request, so its per-request CPU grows with the number of
+        in-flight samples (~28ms each at 512 in flight); aiohttp stays flat.
+        """
+        if self._session is None or self._session.closed:
+            max_conn = self.engine_config.forwarding_inference_max_connections
+            # keepalive_timeout must stay under the router's idle timeout so a
+            # pooled connection is never reused after the server closed it.
+            # Happy Eyeballs is off: a burst of connect timeouts cancels its
+            # sock_connect calls mid-flight, and under uvloop the closed sockets'
+            # descriptors get reused before the loop forgets them ("File
+            # descriptor N is used by transport"), failing unrelated forwards.
+            connector = aiohttp.TCPConnector(limit=max_conn or 0, keepalive_timeout=2, happy_eyeballs_delay=None)
+            self._session = aiohttp.ClientSession(
+                connector=connector,
+                timeout=aiohttp.ClientTimeout(
+                    total=None,
+                    # A saturated router can take tens of seconds to accept;
+                    # that is queueing, not failure.
+                    sock_connect=_ROUTER_CONNECT_TIMEOUT_SECONDS,
+                    sock_read=self.engine_config.forwarding_inference_timeout_sec,
+                ),
+            )
+        return self._session
 
     async def aclose(self) -> None:
-        """Close the persistent httpx client. Called from api.py lifespan shutdown."""
-        await self._http_client.aclose()
+        """Close the shared aiohttp session. Called from api.py lifespan shutdown."""
+        if self._session is not None and not self._session.closed:
+            await self._session.close()
 
     async def _read_proxy_url_from_db(self) -> str | None:
         async with AsyncSession(self.db_engine) as session:
@@ -119,30 +148,31 @@ async def call_and_store_result(
             await session.commit()
 
     async def _forward_with_retry(self, sample_req, model_id: str, *, base_model: str | None) -> types.SampleOutput:
-        # Retry only failures that occur before a request can reach vLLM. Read
-        # and write failures are ambiguous: vLLM may still be executing the
+        # 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
         # request, so retrying would duplicate generation load.
         try:
             try:
                 proxy_url = await self._resolve_proxy_url()
                 return await self._forward(proxy_url, sample_req, model_id, base_model=base_model)
-            except (httpx.ConnectError, httpx.ConnectTimeout) as e:
+            except (aiohttp.ClientConnectorError, aiohttp.ConnectionTimeoutError, TransientInferenceError) as e:
                 logger.warning(
-                    "Connection error talking to %s (%s: %s) — refreshing proxy URL and retrying once",
+                    "Transient error talking to %s (%s: %s) — refreshing proxy URL and retrying once",
                     self._cached_proxy_url,
                     type(e).__name__,
                     e,
                 )
                 proxy_url = await self._resolve_proxy_url(force_refresh=True)
                 return await self._forward(proxy_url, sample_req, model_id, base_model=base_model)
-        except httpx.ReadTimeout as e:
+        except aiohttp.SocketTimeoutError as e:
             # Not retried (see above). Long-context requests routinely exceed the
             # default read deadline, so tell the caller how to raise it. The
             # message is stored in the FutureDB ErrorResponse and shown to clients.
             timeout_sec = self.engine_config.forwarding_inference_timeout_sec
             raise RuntimeError(
                 f"Inference request to {self._cached_proxy_url} timed out after {timeout_sec:g}s waiting for "
-                "a response (httpx.ReadTimeout). The request was not retried because vLLM may still be "
+                "a response (read timeout). The request was not retried because vLLM may still be "
                 "executing it. If requests are expected to take this long (long prompts, large max_tokens, "
                 "or queueing behind other requests), increase the deadline with "
                 "`--forwarding-inference-timeout-sec` (EngineConfig.forwarding_inference_timeout_sec) or "
@@ -195,17 +225,22 @@ async def _forward(
             headers["X-Session-ID"] = session_id
 
         url = f"{proxy_url}/v1/completions"
-        response = await self._http_client.post(url, json=payload, headers=headers)
-        if response.status_code >= 400:
-            raise RuntimeError(f"vLLM /v1/completions returned {response.status_code}: {response.text}")
-        try:
-            result = response.json()
-        except ValueError as e:
-            # vllm-router can return HTML on transient errors even with 2xx status.
-            raise RuntimeError(
-                f"vLLM /v1/completions returned non-JSON ({response.status_code}, "
-                f"content-type={response.headers.get('content-type')!r}): {response.text[:512]}"
-            ) from e
+        async with self._get_session().post(url, json=payload, headers=headers) as response:
+            body = await response.read()
+            if response.status >= 500:
+                raise TransientInferenceError(
+                    f"vLLM /v1/completions returned {response.status}: {body.decode(errors='replace')}"
+                )
+            if response.status >= 400:
+                raise RuntimeError(f"vLLM /v1/completions returned {response.status}: {body.decode(errors='replace')}")
+            try:
+                result = orjson.loads(body)
+            except orjson.JSONDecodeError as e:
+                # vllm-router can return HTML on transient errors even with 2xx status.
+                raise RuntimeError(
+                    f"vLLM /v1/completions returned non-JSON ({response.status}, "
+                    f"content-type={response.headers.get('content-type')!r}): {body[:512].decode(errors='replace')}"
+                ) from e
 
         prompt_logprobs = None
         topk = None

tests/tinker/test_external_future_store.py

@@ -42,6 +42,11 @@ def _sample_input(seq_id: int) -> types.SampleInput:
     )
 
 
+async def _still_connected() -> bool:
+    """Stand-in for ``Request.is_disconnected`` on a client that is still waiting."""
+    return False
+
+
 class _CompletingForwarder:
     def __init__(self, store: ExternalFutureStore):
         self.store = store
@@ -140,6 +145,7 @@ async def test_sustained_model_path_rollouts_training_futures_and_heartbeats(fut
             )
         ),
         headers={},
+        is_disconnected=_still_connected,
     )
 
     async with AsyncSession(engine) as session:
@@ -294,6 +300,7 @@ def serialize_result_in_thread(request_type, result_data):
             )
         ),
         headers={"accept": api.PROTO_CONTENT_TYPE},
+        is_disconnected=_still_connected,
     )
 
     responses = await asyncio.gather(
@@ -643,6 +650,7 @@ async def test_retrieve_future_serializes_in_memory_result_as_proto(future_store
             )
         ),
         headers={"accept": "application/x-protobuf, application/json"},
+        is_disconnected=_still_connected,
     )
     response = await api.retrieve_future(api.RetrieveFutureRequest(request_id=str(request_id)), request)