fix(endpoint-client): retry pre-response connection resets (opt-in)

examples/11_Edge_Agentic_Example/online_edge_full_run.yaml

@@ -106,6 +106,12 @@ settings:
     num_workers: 1
     max_connections: 1
     warmup_connections: 0
+    # 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
 
 endpoint_config:
   endpoints:

src/inference_endpoint/config/templates/concurrency_template_full.yaml

@@ -78,6 +78,7 @@ settings:
     worker_force_kill_timeout: 0.5  # Force kill timeout after graceful wait (seconds)
     insecure: false  # Skip TLS certificate verification
     max_idle_time: 4.0  # Discard connections idle longer than this (seconds)
+    transport_max_retries: 0  # Retries on a pre-response connection reset (0=disabled)
     min_required_connections: -1  # Min connections to initialize (-1=auto, 0=disabled)
     worker_gc_mode: relaxed  # Worker GC strategy | options: disabled, relaxed, system
   drain:

src/inference_endpoint/config/templates/offline_template_full.yaml

@@ -78,6 +78,7 @@ settings:
     worker_force_kill_timeout: 0.5  # Force kill timeout after graceful wait (seconds)
     insecure: false  # Skip TLS certificate verification
     max_idle_time: 4.0  # Discard connections idle longer than this (seconds)
+    transport_max_retries: 0  # Retries on a pre-response connection reset (0=disabled)
     min_required_connections: -1  # Min connections to initialize (-1=auto, 0=disabled)
     worker_gc_mode: relaxed  # Worker GC strategy | options: disabled, relaxed, system
   drain:

src/inference_endpoint/config/templates/online_template_full.yaml

@@ -79,6 +79,7 @@ settings:
     worker_force_kill_timeout: 0.5  # Force kill timeout after graceful wait (seconds)
     insecure: false  # Skip TLS certificate verification
     max_idle_time: 4.0  # Discard connections idle longer than this (seconds)
+    transport_max_retries: 0  # Retries on a pre-response connection reset (0=disabled)
     min_required_connections: -1  # Min connections to initialize (-1=auto, 0=disabled)
     worker_gc_mode: relaxed  # Worker GC strategy | options: disabled, relaxed, system
   drain:

src/inference_endpoint/endpoint_client/config.py

@@ -170,6 +170,23 @@ class HTTPClientConfig(WithUpdatesMixin, BaseModel):
         4.0, description="Discard connections idle longer than this (seconds)"
     )
 
+    # Retry a request when the connection fails BEFORE any response byte is
+    # received (e.g. "connection closed before headers received" — the server
+    # closed an idle keep-alive socket at the instant we wrote the request).
+    # Safe because the server produced no output; the retry re-issues on a fresh
+    # connection. Never fires once headers/chunks have started (would risk
+    # duplicate output). Guards single-stream localhost runs (e.g. llama.cpp with
+    # -np 1) where a stale-socket reset otherwise zeroes an accuracy sample.
+    #
+    # 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(
+        0,
+        ge=0,
+        description="Retries on a pre-response connection reset (0=disabled)",
+    )
+
     # Minimum required connections for http-client to initialize.
     # Will log warning if not enough ephemeral ports are available during warmup.
     #

src/inference_endpoint/endpoint_client/worker.py

@@ -378,14 +378,47 @@ async def _fire_request(self, req: InFlightRequest) -> bool:
             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:
+        """Read response headers, retrying on a pre-response connection reset.
+
+        A connection-level failure at this point means the server closed the
+        socket before sending any response byte (the idle keep-alive race:
+        server drops an idle -np 1 connection just as the client writes on it).
+        No output was produced, so re-issuing the identical request on a fresh
+        connection is safe and idempotent. Returns the HTTP status code, or None
+        if retries are exhausted (an error response has been emitted).
+        """
+        attempt = 0
+        while True:
+            conn = req.connection
+            try:
+                status_code, _ = await conn.protocol.read_headers()
+                return status_code
+            except ConnectionError as e:
+                # Dead connection: discard it (frees the pool slot) and, if
+                # retries remain, re-issue on a fresh connection. Retrying is
+                # only reached here because no headers arrived, so no partial
+                # output was delivered to the accumulator.
+                self._pool.release(conn)
+                if attempt >= self.http_config.transport_max_retries:
+                    await self._handle_error(req.query_id, e)
+                    return None
+                attempt += 1
+                new_conn = await self._pool.acquire()
+                new_conn.protocol.write(req.http_bytes)
+                req.connection = new_conn
+
     @profile
     async def _process_response(self, req: InFlightRequest) -> None:
         """Process response for a fired request."""
-        conn = req.connection
-
         try:
-            # Await headers and handle error status
-            status_code, _ = await conn.protocol.read_headers()
+            # Await headers, retrying on a pre-response connection reset.
+            status_code = await self._read_headers_with_retry(req)
+            if status_code is None:
+                # Retries exhausted; error already emitted.
+                return
+
+            conn = req.connection
             if status_code != 200:
                 error_body = await conn.protocol.read_body()
                 self._pool.release(conn)
@@ -406,8 +439,9 @@ async def _process_response(self, req: InFlightRequest) -> None:
             logger.warning(f"Request {req.query_id} failed: {type(e).__name__}: {e}")
 
         finally:
-            # Release connection back to pool if not already
-            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)
 
             # Clean up task reference
             current_task = asyncio.current_task()

tests/unit/endpoint_client/test_worker_retry.py

@@ -0,0 +1,192 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the worker's pre-response connection-reset retry.
+
+Covers ``Worker._read_headers_with_retry`` — the guard that re-issues a request
+on a fresh connection when the server closes the socket before sending any
+response byte (the idle keep-alive race that otherwise zeroes accuracy samples
+on single-stream localhost servers such as llama.cpp ``-np 1``).
+"""
+
+from types import SimpleNamespace
+
+import pytest
+from inference_endpoint.endpoint_client.config import HTTPClientConfig
+from inference_endpoint.endpoint_client.http import InFlightRequest
+from inference_endpoint.endpoint_client.worker import Worker
+
+
+@pytest.mark.unit
+def test_transport_retries_default_is_opt_in():
+    """Default is 0 (fail-fast) so the retry never changes behavior unless a
+    config opts in explicitly (e.g. the edge-agentic accuracy reference)."""
+    assert HTTPClientConfig().transport_max_retries == 0
+
+
+class _FakeProtocol:
+    """Protocol whose ``read_headers`` replays a scripted result sequence."""
+
+    def __init__(self, results):
+        # results: list where each item is either the sentinel "reset" (raise
+        # ConnectionResetError) or an int status code to return.
+        self._results = list(results)
+        self.written: list[bytes] = []
+
+    async def read_headers(self):
+        outcome = self._results.pop(0)
+        if outcome == "reset":
+            raise ConnectionResetError("Connection closed before headers received")
+        return outcome, {}
+
+    def write(self, data: bytes) -> None:
+        self.written.append(data)
+
+
+class _FakeConn:
+    def __init__(self, protocol: _FakeProtocol):
+        self.protocol = protocol
+
+
+class _FakePool:
+    """Records releases and hands out queued connections on ``acquire``."""
+
+    def __init__(self, acquire_queue: list[_FakeConn]):
+        self._acquire_queue = list(acquire_queue)
+        self.released: list[_FakeConn] = []
+        self.acquired: list[_FakeConn] = []
+
+    def release(self, conn) -> None:
+        self.released.append(conn)
+
+    async def acquire(self) -> _FakeConn:
+        conn = self._acquire_queue.pop(0)
+        self.acquired.append(conn)
+        return conn
+
+
+def _make_worker(pool: _FakePool, max_retries: int):
+    """Build a Worker with only the attributes the retry path touches."""
+    worker = Worker.__new__(Worker)
+    worker._pool = pool  # type: ignore[assignment]
+    worker.http_config = SimpleNamespace(  # type: ignore[assignment]
+        transport_max_retries=max_retries
+    )
+    worker._handle_error_calls = []  # type: ignore[attr-defined]
+
+    async def _record_error(query_id, error):
+        worker._handle_error_calls.append((query_id, error))
+
+    worker._handle_error = _record_error  # type: ignore[method-assign]
+    return worker
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_retry_reissues_on_fresh_connection_after_reset():
+    """A pre-response reset is retried on a fresh connection and succeeds."""
+    dead = _FakeConn(_FakeProtocol(["reset"]))
+    fresh = _FakeConn(_FakeProtocol([200]))
+    pool = _FakePool(acquire_queue=[fresh])
+    worker = _make_worker(pool, max_retries=2)
+
+    req = InFlightRequest(
+        query_id="q1",
+        http_bytes=b"POST /v1/chat/completions HTTP/1.1\r\n\r\n",
+        is_streaming=True,
+        connection=dead,
+    )
+
+    status = await worker._read_headers_with_retry(req)
+
+    assert status == 200
+    assert req.connection is fresh
+    # Dead connection discarded exactly once; request re-written on the fresh one.
+    assert pool.released == [dead]
+    assert fresh.protocol.written == [req.http_bytes]
+    assert worker._handle_error_calls == []
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_retry_exhausted_emits_error_and_returns_none():
+    """When every attempt resets, retries are exhausted and an error is emitted."""
+    dead = _FakeConn(_FakeProtocol(["reset"]))
+    also_dead = _FakeConn(_FakeProtocol(["reset"]))
+    pool = _FakePool(acquire_queue=[also_dead])
+    worker = _make_worker(pool, max_retries=1)
+
+    req = InFlightRequest(
+        query_id="q2",
+        http_bytes=b"POST /x HTTP/1.1\r\n\r\n",
+        is_streaming=False,
+        connection=dead,
+    )
+
+    status = await worker._read_headers_with_retry(req)
+
+    assert status is None
+    # Both connections were tried and discarded.
+    assert pool.released == [dead, also_dead]
+    assert len(worker._handle_error_calls) == 1
+    assert worker._handle_error_calls[0][0] == "q2"
+    assert isinstance(worker._handle_error_calls[0][1], ConnectionResetError)
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_retries_disabled_emits_error_without_reacquiring():
+    """transport_max_retries=0 => no re-issue; the reset surfaces immediately."""
+    dead = _FakeConn(_FakeProtocol(["reset"]))
+    pool = _FakePool(acquire_queue=[])  # acquire must never be called
+    worker = _make_worker(pool, max_retries=0)
+
+    req = InFlightRequest(
+        query_id="q3",
+        http_bytes=b"POST /x HTTP/1.1\r\n\r\n",
+        is_streaming=False,
+        connection=dead,
+    )
+
+    status = await worker._read_headers_with_retry(req)
+
+    assert status is None
+    assert pool.acquired == []
+    assert pool.released == [dead]
+    assert len(worker._handle_error_calls) == 1
+
+
+@pytest.mark.unit
+@pytest.mark.asyncio
+async def test_headers_ok_first_try_never_touches_pool():
+    """No reset => no release/acquire churn, status returned directly."""
+    conn = _FakeConn(_FakeProtocol([200]))
+    pool = _FakePool(acquire_queue=[])
+    worker = _make_worker(pool, max_retries=2)
+
+    req = InFlightRequest(
+        query_id="q4",
+        http_bytes=b"POST /x HTTP/1.1\r\n\r\n",
+        is_streaming=False,
+        connection=conn,
+    )
+
+    status = await worker._read_headers_with_retry(req)
+
+    assert status == 200
+    assert req.connection is conn
+    assert pool.acquired == []
+    assert pool.released == []
+    assert worker._handle_error_calls == []