fix(server_utils): enable TCP keepalive on global aiohttp connector

[martinagvilas]

Symptom

Long-running Gym benchmarks (e.g. AA-Omniscience, MMLU-Pro) occasionally freeze with a handful of rollouts stuck "in flight" forever, until the Slurm job hits its time limit.

Root cause

The global aiohttp client uses a default TCPConnector with no TCP keepalive enabled. When an upstream socket is silently dropped, it is not noticed, so the request blocks forever.

Not sure whether this was intentional. The connector's existing keepalive_timeout=15.0 is not TCP keepalive, it only recycles unused connections in the pool, so it doesn't help with stuck in-flight requests.

Fix

_TCPKeepAliveConnector: a TCPConnector subclass that turns on SO_KEEPALIVE on every new connection with tight timings (idle 60s, probe every 10s, give up after 3 probes). Once the kernel drops the socket, aiohttp raises ClientOSError, which the existing request() retry loop already handles.

Implemented via _wrap_create_connection rather than TCPConnector(socket_options=…) so it works on aiohttp<3.11 too. Linux-only sockopts are guarded so the code also runs on macOS/BSD.

Happy to bump or lower the specific timings (60s / 10s / 3 probes).

Tests

Three unit tests cover: happy path (all keepalive sockopts set), non-Linux platform (only SO_KEEPALIVE set, no crash), and fail-loud when the transport has no underlying socket.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[cmunley1]

/claude review

[bxyu-nvidia]

can we parameterize this via GlobalAIOHTTPAsyncClientConfig? and what are the defaults for these values if we don't specify them, can we default to those same defaults or at least display them in a comment here?

[martinagvilas]

Good idea. I moved the three TCP keepalive values (TCP_KEEPIDLE, TCP_KEEPINTVL, TCP_KEEPCNT) to GlobalAIOHTTPAsyncClientConfig fields so they're overridable per deployment. Linux kernel defaults are 7200s / 75s / 9 probes, so ~2h15m before TCP gives up. I think that's longer than most benchmarks need, so I used shorter times.

[claude]

SHIP WITH CARE — the fix is correct and well-targeted; one failure-mode asymmetry is worth reconsidering before merge.

The diagnosis is right: keepalive_timeout=15.0 only recycles idle pool connections, it does nothing for a silently-dropped in-flight socket. Enabling SO_KEEPALIVE + TCP_KEEPIDLE/INTVL/CNT turns a forever-hang into a ClientOSError that the existing request() retry loop already handles (server_utils.py:215). The ~90s detection budget is reasonable for long rollouts. Linux-only sockopts are correctly guarded with getattr, and the three unit tests cover happy path, non-Linux, and no-socket. Good change.

Two things (both inline):

1. RISK (server_utils.py:113): the hard raise RuntimeError when sock is None inverts the failure mode — a keepalive optimization becomes a total-outage switch, because the condition is deterministic per transport type and would fail every request, not one. Degrade gracefully (warn once, skip sockopts) instead of raising in the hot path. The downside of losing keepalive is strictly milder than the problem being fixed.

2. NOTE (server_utils.py:109): the _wrap_create_connection override cites "aiohttp<3.11 compat", but pyproject floors aiohttp at >=3.14.1. The public TCPConnector(socket_options=...) API (3.11+) is available on every supported version and avoids overriding a private method. Works as-is; just no longer needed.

Nothing here corrupts scores or breaks the public server API — it's transport-layer reliability. Fix #1 and it's a clean ship.

[claude]

RISK — fail-hard on a keepalive optimization inverts the intended failure mode.

WHAT BREAKS: _wrap_create_connection runs on every new connection. If get_extra_info("socket") ever returns None for a transport type in production, this raise propagates out of client.request(), gets caught by the generic except Exception in request() (line 227), retries MAX_NUM_TRIES, then re-raises — permanently. The condition is deterministic per transport type, so it wouldn't fail one request, it would fail all of them.

BLAST RADIUS: a feature added to prevent stalled rollouts becomes a total-outage switch. Losing keepalive is a degraded state (back to today's behavior); losing all HTTP is an outage. The downside of the guard is strictly worse than the problem it guards against.

FIX: degrade gracefully instead of raising — log/warn once and return transport, protocol without setting sockopts. Keep the loud raise only in the test (assert the warning), not in the hot path. On Linux+asyncio/uvloop with TCP/SSL transports socket is reliably present, so this is defense-in-depth, but the asymmetry (trivial fix, catastrophic failure mode) argues for not raising.

[martinagvilas]

Switched the raise RuntimeError in _wrap_create_connection to a log-once warning that returns the transport unchanged. If a transport ever exposes no underlying socket, we now just skip keepalive on that connection instead of taking down the whole client through the retry loop. Test updated to assert graceful degradation and a single warning per connector.