[tinker] Forward samples with aiohttp instead of httpx

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/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_inference_forwarding_config.py

@@ -1,12 +1,14 @@
 import argparse
-from unittest.mock import AsyncMock, call, patch
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, call
 
-import httpx
+import aiohttp
 import pytest
 
 from skyrl.tinker.config import EngineConfig, add_model
 from skyrl.tinker.extra.skyrl_train_inference_forwarding import (
     SkyRLTrainInferenceForwardingClient,
+    TransientInferenceError,
 )
 
 
@@ -21,20 +23,37 @@ def test_forwarding_timeout_reads_environment(monkeypatch) -> None:
     assert config.forwarding_inference_timeout_sec == 1800.0
 
 
-def test_forwarding_client_uses_configured_timeout() -> None:
+@pytest.mark.asyncio
+async def test_forwarding_client_uses_configured_timeout_and_connection_limit() -> None:
     config = EngineConfig(
         base_model="test-model",
         forwarding_inference_timeout_sec=1800.0,
+        forwarding_inference_max_connections=64,
     )
+    client = SkyRLTrainInferenceForwardingClient(config, db_engine=None)
+    try:
+        session = client._get_session()
+        assert session.timeout.sock_connect == 60.0
+        assert session.timeout.sock_read == 1800.0
+        # No overall deadline: a request may wait in the connector queue for
+        # as long as the engine takes to get to it.
+        assert session.timeout.total is None
+        assert session.connector.limit == 64
+    finally:
+        await client.aclose()
+
+
+@pytest.mark.asyncio
+async def test_forwarding_client_default_connection_limit_is_unlimited() -> None:
+    client = SkyRLTrainInferenceForwardingClient(EngineConfig(base_model="test-model"), db_engine=None)
+    try:
+        assert client._get_session().connector.limit == 0
+    finally:
+        await client.aclose()
 
-    with patch("skyrl.tinker.extra.skyrl_train_inference_forwarding.httpx.AsyncClient") as async_client:
-        SkyRLTrainInferenceForwardingClient(config, db_engine=None)
 
-    timeout = async_client.call_args.kwargs["timeout"]
-    assert timeout.connect == 10.0
-    assert timeout.read == 1800.0
-    assert timeout.write == 300.0
-    assert timeout.pool == 300.0
+def _connect_error(message: str) -> aiohttp.ClientConnectorError:
+    return aiohttp.ClientConnectorError(SimpleNamespace(ssl=None, host="inference", port=8000), OSError(message))
 
 
 @pytest.mark.asyncio
@@ -43,7 +62,7 @@ async def test_forwarding_retries_connection_failure() -> None:
     client._cached_proxy_url = "http://old"
     client._resolve_proxy_url = AsyncMock(side_effect=["http://old", "http://new"])
     expected = object()
-    client._forward = AsyncMock(side_effect=[httpx.ConnectError("unreachable"), expected])
+    client._forward = AsyncMock(side_effect=[_connect_error("unreachable"), expected])
 
     result = await client._forward_with_retry(object(), "model", base_model=None)
 
@@ -52,19 +71,46 @@ async def test_forwarding_retries_connection_failure() -> None:
     assert client._forward.await_count == 2
 
 
+@pytest.mark.asyncio
+async def test_forwarding_retries_transient_5xx_once() -> None:
+    client = object.__new__(SkyRLTrainInferenceForwardingClient)
+    client._cached_proxy_url = "http://old"
+    client._resolve_proxy_url = AsyncMock(side_effect=["http://old", "http://new"])
+    expected = object()
+    client._forward = AsyncMock(side_effect=[TransientInferenceError("503 from router"), expected])
+
+    result = await client._forward_with_retry(object(), "model", base_model=None)
+
+    assert result is expected
+    assert client._forward.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_forwarding_does_not_retry_4xx() -> None:
+    client = object.__new__(SkyRLTrainInferenceForwardingClient)
+    client._cached_proxy_url = "http://inference"
+    client._resolve_proxy_url = AsyncMock(return_value="http://inference")
+    client._forward = AsyncMock(side_effect=RuntimeError("vLLM /v1/completions returned 400: bad request"))
+
+    with pytest.raises(RuntimeError, match="returned 400"):
+        await client._forward_with_retry(object(), "model", base_model=None)
+
+    client._forward.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_forwarding_does_not_retry_read_timeout() -> None:
     client = object.__new__(SkyRLTrainInferenceForwardingClient)
     client.engine_config = EngineConfig(base_model="test-model", forwarding_inference_timeout_sec=123.0)
     client._cached_proxy_url = "http://inference"
     client._resolve_proxy_url = AsyncMock(return_value="http://inference")
-    client._forward = AsyncMock(side_effect=httpx.ReadTimeout("slow response"))
+    client._forward = AsyncMock(side_effect=aiohttp.SocketTimeoutError("slow response"))
 
     with pytest.raises(RuntimeError) as exc_info:
         await client._forward_with_retry(object(), "model", base_model=None)
 
     message = str(exc_info.value)
-    assert isinstance(exc_info.value.__cause__, httpx.ReadTimeout)
+    assert isinstance(exc_info.value.__cause__, aiohttp.SocketTimeoutError)
     assert "http://inference" in message
     assert "timed out after 123s" in message
     client._resolve_proxy_url.assert_awaited_once_with()