diff --git a/skyrl/tinker/config.py b/skyrl/tinker/config.py index cd8d5c4a23..a88c005f30 100644 --- a/skyrl/tinker/config.py +++ b/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, diff --git a/skyrl/tinker/extra/skyrl_train_inference_forwarding.py b/skyrl/tinker/extra/skyrl_train_inference_forwarding.py index d617d9a7a6..88e63a6776 100644 --- a/skyrl/tinker/extra/skyrl_train_inference_forwarding.py +++ b/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 diff --git a/tests/tinker/test_inference_forwarding_config.py b/tests/tinker/test_inference_forwarding_config.py index 091a49efb3..1987fb1deb 100644 --- a/tests/tinker/test_inference_forwarding_config.py +++ b/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()