[tinker] Forward samples with aiohttp instead of httpx

[avigyabb]

Stack 2/7.

Why. httpcore's connection pool rescans every pooled connection and queued request on each event, so the forwarding client's per-request CPU grows with the number of in-flight samples: 14 ms at 512, 27 ms at 2048, 38 ms at 4096 (standalone benchmark against an instant fake router; the profile shows 27.5M is_idle calls for 512 requests). At 2048 in flight the API server needed 119 s to forward 2048 instant requests. aiohttp's connector is flat at 0.34 ms.

Change. SkyRLTrainInferenceForwardingClient uses one aiohttp session:

• connector limit = forwarding_inference_max_connections (0 = unlimited); sock_read = forwarding_inference_timeout_sec; no total deadline, so requests queued behind the engine never hit the old 300 s pool timeout;
• 60 s connect timeout (a saturated router takes tens of seconds to accept; that is queueing, not failure) and Happy Eyeballs off: a wave of cancelled connects under uvloop produced ~800 File descriptor N is used by transport errors from aiohappyeyeballs' sock_connect, failing unrelated forwards;
• connect-phase errors and 5xx rejections (TransientInferenceError) are retried once after refreshing the proxy URL; read failures stay final (fix(tinker): bound inference forwarding retries #2118 stance) since vLLM may still be executing the request;
• forwarding_inference_timeout_sec default 300 s -> 2048 s: with unlimited connections a large burst waits inside vLLM's queue and 128x128 bursts exceed 300 s there.

Results (same harness and SDK-shaped client, upstream main vs this stack): 2048 small results unlimited connections 71.6 s / 29 per s -> 2.7 s / 747 per s; 32768 with 5 s engine queueing 218 s with 530 failures -> 89 s, 0 failures.

Note: unlimited outbound connections cannot exceed ~28k (one source IP's ephemeral ports); set --forwarding-inference-max-connections near engine capacity for very large bursts.

Stack (each PR retargets to main as the one below merges)

1. [tinker] Route uvicorn's access log to a plain handler instead of Rich #2160
2. [tinker] Forward samples with aiohttp instead of httpx #2161
3. [tinker] Keep an undelivered sample result alive for the SDK's retry #2162
4. [tinker] Survive completion bursts at the socket layer (accept backlog, keep-alive) #2163
5. [tinker] Encode forwarded sample results to proto once and serve them as-is #2164
6. [tinker] Decode vLLM completion bodies straight into numpy with pysimdjson #2165
7. [tinker] Load harness for the API server's sampling path at 131k concurrency #2166

🤖 Generated with Claude Code

---

Note

Medium Risk
Touches hot-path inference forwarding, connection/retry semantics, and future eviction timing under load; mis-tuned timeouts or disconnect handling could cause duplicate forwards or 404s on SDK retries.

Overview
Replaces httpx with aiohttp in SkyRLTrainInferenceForwardingClient so high fan-out rollouts don’t pay growing per-request pool CPU; the shared session uses connector limits, a 60s connect timeout, no total deadline (queueing waits in the connector), Happy Eyeballs off under uvloop, and orjson for response bodies. Retries refresh the proxy URL once on connect failures and on router 5xx (TransientInferenceError); read timeouts stay non-retryable.

forwarding_inference_timeout_sec default rises 300 → 2048 to cover long vLLM queue waits on large bursts.

Uvicorn gets timeout_keep_alive=75 and backlog=SKYRL_HTTP_CONNECTION_LIMIT so SDK polls survive completion bursts without mass reconnect/refused accepts.

retrieve_future only calls mark_retrieved when the client is still connected, so abandoned 45s SDK polls don’t start the short TTL on undelivered results; _RETRIEVED_TTL_SECONDS is 120 → 300 to cover SDK retry/backoff. Integration tests cover the abandoned-poll → retry path over a real socket.

^{Reviewed by Cursor Bugbot for commit 3bdbb3d. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request replaces httpx with aiohttp and orjson in the inference forwarding client to improve performance and connection handling under high concurrency, and increases the default forwarding timeout to 2048 seconds. The reviewer identified a critical issue where the code attempts to catch non-existent exceptions (aiohttp.ConnectionTimeoutError and aiohttp.SocketTimeoutError), which would lead to runtime AttributeErrors. To resolve this, the reviewer suggested defining custom exceptions and wrapping the connection establishment and response reading in separate try-except blocks to properly distinguish between connection and read timeouts (as both raise asyncio.TimeoutError in aiohttp), as well as updating the corresponding tests.

[gemini-code-assist]

The aiohttp library does not define ConnectionTimeoutError or SocketTimeoutError exceptions. Referencing them will raise an AttributeError at runtime.

To correctly handle and distinguish connection timeouts from read timeouts, we can define custom ConnectionTimeoutError and ReadTimeoutError exceptions.

class TransientInferenceError(RuntimeError):
    """A 5xx from vllm-router/vLLM: the request was rejected, not executed, so it is safe to retry."""


class ConnectionTimeoutError(RuntimeError):
    """Connection timed out."""


class ReadTimeoutError(RuntimeError):
    """Read timed out."""


_ROUTER_CONNECT_TIMEOUT_SECONDS = 60.0

[gemini-code-assist]

Catch the custom ConnectionTimeoutError and ReadTimeoutError exceptions instead of the non-existent aiohttp.ConnectionTimeoutError and aiohttp.SocketTimeoutError.

Suggested change

	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:
	except (aiohttp.ClientConnectorError, ConnectionTimeoutError, TransientInferenceError) as e:
	logger.warning(
	"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 ReadTimeoutError as e:

[gemini-code-assist]

In aiohttp, both connection timeouts and read timeouts raise asyncio.TimeoutError. To distinguish them, we can wrap the connection establishment (async with self._get_session().post(...)) and the response reading (await response.read()) in separate try...except blocks, raising the custom ConnectionTimeoutError and ReadTimeoutError respectively.

        try:
            async with self._get_session().post(url, json=payload, headers=headers) as response:
                try:
                    body = await response.read()
                except asyncio.TimeoutError as e:
                    raise ReadTimeoutError("Read timeout") from e
                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
        except asyncio.TimeoutError as e:
            raise ConnectionTimeoutError("Connection timeout") from e

[gemini-code-assist]

Import the custom ConnectionTimeoutError and ReadTimeoutError exceptions for testing.

Suggested change

	from skyrl.tinker.extra.skyrl_train_inference_forwarding import (
	SkyRLTrainInferenceForwardingClient,
	TransientInferenceError,
	)
	from skyrl.tinker.extra.skyrl_train_inference_forwarding import (
	SkyRLTrainInferenceForwardingClient,
	TransientInferenceError,
	ConnectionTimeoutError,
	ReadTimeoutError,
	)

[gemini-code-assist]

Update the test to use the custom ReadTimeoutError instead of the non-existent aiohttp.SocketTimeoutError.

Suggested change

	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)
	client._forward = AsyncMock(side_effect=ReadTimeoutError("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__, ReadTimeoutError)

[greptile-apps]

Confidence Score: 4/5

The PR should not merge until the blanket 5xx retry is narrowed or made idempotent, because a server error can occur after inference has already been accepted.

The new response handling resubmits every 5xx completion without a deduplication boundary, allowing ambiguous server failures to trigger duplicate generation work.

Files Needing Attention: skyrl/tinker/extra/skyrl_train_inference_forwarding.py

Important Files Changed

Filename	Overview
skyrl/tinker/extra/skyrl_train_inference_forwarding.py	Migrates forwarding to aiohttp, but the blanket retry of every 5xx can duplicate an already accepted inference request.
skyrl/tinker/config.py	Raises the default forwarding read timeout and documents why large queued bursts require it.
tests/tinker/test_inference_forwarding_config.py	Updates timeout, connector-limit, and retry tests, though the mocked 5xx test assumes all server errors are safe to retry.

Sequence Diagram

sequenceDiagram
    participant API
    participant Router
    participant vLLM
    API->>Router: POST /v1/completions
    Router->>vLLM: Forward completion
    vLLM-->>Router: 5xx after accepting request
    Router-->>API: 5xx
    Note over API: Classify every 5xx as transient
    API->>Router: Retry same completion
    Router->>vLLM: Execute completion again

Loading

_{Reviews (1): Last reviewed commit: "[tinker] Forward samples with aiohttp in..." | Re-trigger Greptile}

[greptile-apps]

5xx retries duplicate inference

When vLLM or the router returns a 5xx after accepting a completion, _forward classifies it as safe to retry and submits the same non-idempotent request again, causing duplicate generation work and additional engine load.