fix(endpoint-client): retry pre-response connection resets (opt-in)#404
fix(endpoint-client): retry pre-response connection resets (opt-in)#404Palanivelg wants to merge 2 commits into
Conversation
Single-slot servers that recycle idle keep-alive sockets between requests can close a socket at the instant the client reuses it. The write then races the close and raises a connection reset before any response byte arrives, which previously failed the request outright -- and on an accuracy run, silently zeroed the affected sample. Add Worker._read_headers_with_retry: on a connection failure before any response byte, discard the dead connection and re-issue the identical request on a fresh one. It never fires once headers/chunks have started, so it cannot duplicate output. Gated by a new HTTPClientConfig field transport_max_retries (default 0 = prior fail-fast behavior, so every existing use case is unchanged); the edge-agentic accuracy reference config opts in explicitly. Adds unit tests covering retry-success, retries-exhausted, retries-disabled, the no-reset fast path, and the opt-in default.
|
MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅ |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a retry mechanism for HTTP requests that fail due to pre-response connection resets, which commonly occurs when a server closes an idle keep-alive socket just as the client sends a request. This behavior is controlled by a new transport_max_retries configuration parameter. The feedback highlights a potential connection leak in the retry loop of _read_headers_with_retry if writing to a newly acquired connection fails before it is assigned to the request, and suggests assigning the connection immediately after acquisition to ensure it is properly released in the finally block.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| new_conn.protocol.write(req.http_bytes) | ||
| req.connection = new_conn | ||
|
|
There was a problem hiding this comment.
If new_conn.protocol.write(req.http_bytes) raises an exception (e.g., due to a closed transport or other unexpected errors), the newly acquired connection new_conn will be leaked. This is because req.connection has not yet been updated to new_conn, so the finally block in _process_response will only release the old, already-released connection.\n\nTo prevent this resource leak, assign req.connection = new_conn immediately after acquiring it, before performing the write operation.
new_conn = await self._pool.acquire()\n req.connection = new_conn\n new_conn.protocol.write(req.http_bytes)Assign req.connection to the freshly acquired connection before writing the retried request, so a write failure leaves the request owning a releasable connection (the caller's finally-block releases it) instead of leaking it. Adds a regression test for the failing-write path.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #404 +/- ##
=======================================
Coverage ? 80.04%
=======================================
Files ? 129
Lines ? 17245
Branches ? 0
=======================================
Hits ? 13804
Misses ? 3441
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| # Single-slot edge servers (e.g. llama.cpp -np 1) close idle keep-alive | ||
| # sockets between samples; a reused-then-reset socket would otherwise zero | ||
| # the affected accuracy sample. Retry the request on a fresh connection when | ||
| # the reset happens before any response byte (safe: the server produced no | ||
| # output). Off by default globally; enabled here for the gated run. | ||
| transport_max_retries: 2 |
There was a problem hiding this comment.
@viraatc can you help review this PR? In general I want to avoid adding more tuning flags to our http client if possible
arekay-nv
left a comment
arekay-nv
left a comment
There was a problem hiding this comment.
Review Council — Multi-AI Code Review
Reviewed by: Codex (gpt-5.5, xhigh) + Claude + manual trace | Depth: thorough
Small diff, but two design decisions land on the reproducibility / perf-masking concerns raised for this PR. 6 issues posted inline.
🔴 Must Fix (high)
| File:Line | Category | Summary |
|---|---|---|
worker.py:396 |
observability | Successful retry is invisible (no log/metric/counter/flag/event) → server resets under load masked as clean successes instead of surfaced to tune connection settings |
worker.py:381 |
performance | Retried-sample latency (dead-socket wait + reacquire + re-write + full 2nd server pass) is folded into the sample and attributed to the server with no marker; inflates p99/max on exactly the single-slot samples this targets |
config.py:184 |
reproducibility | transport_max_retries not pinned by check_config_lock and unbounded (ge=0, no le=); recovered resets can silently satisfy no_dropped_turns; reference config sets it to 2 |
🟡 Should Fix (medium)
| File:Line | Category | Summary |
|---|---|---|
worker.py:397 |
correctness | Retries on a partial-response close, not just zero-byte → possible duplicate server-side generation; "safe and idempotent" overstated (Codex + Claude) |
worker.py:447 |
concurrency | release not idempotent for a dead conn (never clears in_use) → double-release fires spurious _notify_waiter; comment inaccurate |
test_worker_retry.py:124 |
testing | Tests never drive _process_response, don't model in_use, no multi-retry-then-success, no observability assertion |
Note: the earlier agent comment on worker.py (write-before-assign leak) is resolved by commit b9ca904.
| conn = req.connection | ||
| try: | ||
| status_code, _ = await conn.protocol.read_headers() | ||
| return status_code |
There was a problem hiding this comment.
[Council: Claude + manual trace] high — perf-regression masking
A successful retry returns the status here with zero trace — no log, no metric, no counter, no per-sample flag, no event. Only the exhausted path (_handle_error, ~L404) and the disabled path record anything. So when a server resets connections under load — a "SUT can't keep up" signal you'd use to tune max_connections/max_idle_time/warmup_connections/server keep-alive — the retry converts it into a clean success and the signal never reaches the metrics/report.
Suggest: logger.warning(...) on each recovery + an aggregatable counter (e.g. transport_retries_total) + a per-sample retry flag on QueryResult, so the Report can surface how many samples needed re-issue.
| logger.error(f"Request {req.query_id} failed: {type(e).__name__}: {e}") | ||
| return False | ||
|
|
||
| async def _read_headers_with_retry(self, req: InFlightRequest) -> int | None: |
There was a problem hiding this comment.
[Council: Claude + manual trace] high — latency pollution attributed to the server
This whole method runs inside the timed _process_response, so a retried sample's elapsed window now folds in: waiting for connection_lost on the dead socket (server FIN/RST timing) + _pool.acquire() of a fresh conn + re-write() + a full second server processing pass — with no marker distinguishing it. On single-slot accuracy runs (the exact scenario this targets) resets are the common case, so this inflates p99/max for precisely the samples the feature exists to handle.
Suggest flagging retried samples in QueryResult so the report can segregate/exclude them instead of blending polluted latency into the server's distribution.
| # Opt-in: default 0 preserves the prior fail-fast behavior for every use | ||
| # case. Configs that need it (e.g. the edge-agentic accuracy reference) | ||
| # enable it explicitly via `client.transport_max_retries`. | ||
| transport_max_retries: int = Field( |
There was a problem hiding this comment.
[Council: Claude + manual trace] high — reproducibility / compliance hole on this knob
compliance/checker.py:check_config_lockpins seed + single-stream (num_workers/max_connections/target_concurrency) but never inspectstransport_max_retries, even though the resolved config carries it atsettings.client.transport_max_retriesand the edge-agentic reference (online_edge_full_run.yaml:114) sets it to2. The config-lock gate therefore passes for any value, it isn't surfaced as a manual attestation, and a retry-recovered reset can let theno_dropped_turnsgate pass on a run where a turn was effectively dropped-and-re-issued.- The field is unbounded (
ge=0, nole=): against a genuinely-down server every attempt re-acquires/re-writes/re-raises, turning fast-fail into a long connection-churn spin that masks a hard SUT failure and multiplies load.
Suggest: pin transport_max_retries in check_config_lock to the ruleset value (0 for perf, the disclosed value for gated accuracy), add le= (e.g. 5), and record recovered resets so "reproducible + 0 dropped turns" can't be satisfied silently. (Dovetails with the maintainer note about limiting tuning flags.)
| try: | ||
| status_code, _ = await conn.protocol.read_headers() | ||
| return status_code | ||
| except ConnectionError as e: |
There was a problem hiding this comment.
[Council: Codex + Claude + manual trace] medium — retries on a partial-response close, not just a true zero-byte close
connection_lost raises ConnectionResetError whenever headers aren't complete (http.py ~L216), so a peer that emitted partial status-line/header bytes and then dropped still lands in this except and re-issues. Client output isn't duplicated (the accumulator only runs after a 200), but the server may have already received and processed the request → duplicate generation work, i.e. extra invisible load that worsens the very overload that caused the reset. The docstring's "no output was produced… safe and idempotent" and the comment "no headers arrived" overstate this — it's client-output-idempotent, not server-side-idempotent.
Suggest tracking whether any response bytes were received and only retrying true zero-byte closes; scope the docstring claim to client output.
| self._pool.release(conn) | ||
| # Release the current connection (body handlers release early; this | ||
| # is idempotent and also covers the swapped connection after a retry) | ||
| self._pool.release(req.connection) |
There was a problem hiding this comment.
[Council: Claude + manual trace] medium — "idempotent" is inaccurate for a dead connection
ConnectionPool.release (http.py:586) only early-returns on not conn.in_use, and its dead branch closes the transport without clearing in_use. In the retry-exhausted and retries-disabled paths the last (dead) connection is released once inside _read_headers_with_retry (~L402) and again here, so the second call re-enters the dead branch and fires a spurious _notify_waiter(). Impact is low in single-slot mode (a spurious wakeup that re-waits), but the comment will mislead maintainers.
Suggest either resetting in_use in the dead branch of release, or dropping the idempotency claim and guarding the double-release explicitly.
|
|
||
| @pytest.mark.unit | ||
| @pytest.mark.asyncio | ||
| async def test_retry_exhausted_emits_error_and_returns_none(): |
There was a problem hiding this comment.
[Council: Claude] medium — test gaps on this PR's risk areas
Every test calls _read_headers_with_retry directly and never drives _process_response, so the finally: self._pool.release(req.connection) double-release path (worker.py L447) is never exercised. _FakeConn/_FakePool don't model in_use, so they couldn't detect the dead-conn idempotency defect even if driven. There's no multi-retry-then-success case (e.g. ["reset","reset",200] with max_retries=2), and nothing asserts observability of a successful retry (there's nothing to assert — itself the gap).
Suggest an integration-style test that drives _process_response end-to-end with an in_use-tracking fake pool.
arekay-nv
left a comment
arekay-nv
left a comment
There was a problem hiding this comment.
Can we root cause this more accurately. Ideally we would like to tune the other settings and not have silent retries impacting the performance.
4 participants
Summary
HTTPClientConfig.transport_max_retries— default0= prior fail-fast behavior, so every existing use case is unchanged. The edge-agentic accuracy reference config opts in explicitly.Changes
endpoint_client/worker.py:_read_headers_with_retryhelperendpoint_client/config.py:transport_max_retriesfield (default 0)examples/11_Edge_Agentic_Example/online_edge_full_run.yaml: opts in_fullconfig templates (schema-driven)tests/unit/endpoint_client/test_worker_retry.py: unit testsTest plan
pytest tests/unit/endpoint_client/test_worker_retry.py— all passpre-commit runon all changed files (ruff, ruff-format, mypy, template regen, license) — all pass