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/async_utils/services/metrics_aggregator/aggregator.py

@@ -70,6 +70,11 @@ class MetricCounterKey(str, Enum):
     # tracked row still exists when the aggregator sees the ERROR.
     TRACKED_SAMPLES_FAILED = "tracked_samples_failed"
     TRACKED_DURATION_NS = "tracked_duration_ns"
+    # Total pre-response connection-reset re-issues across the run (opt-in
+    # transport_max_retries). Non-zero means the SUT reset idle connections
+    # that the client transparently retried — a "can't keep up" signal that
+    # would otherwise vanish once the retry converts the reset into a success.
+    TRANSPORT_RETRIES_TOTAL = "transport_retries_total"
     # Legacy MLPerf LoadGen Server "completed" window (final_query_all_samples_done_time).
     LEGACY_LOADGEN_WINDOW_DURATION_NS = "legacy_loadgen_window_duration_ns"
     # Total wall-clock duration since session start. Updated on every event as
@@ -368,6 +373,13 @@ async def process(self, records: list[EventRecord]) -> None:
                 logger.debug("Error event: %s", record)
                 continue
 
+            # --- Transport retry events ---
+            # A SampleEventType, but not a tracked lifecycle event: it only
+            # bumps the diagnostic recovery counter and never touches the table.
+            if ev == SampleEventType.TRANSPORT_RETRY:
+                registry.increment(MetricCounterKey.TRANSPORT_RETRIES_TOTAL.value)
+                continue
+
             # --- Sample events ---
             if (
                 not isinstance(ev, SampleEventType)

src/inference_endpoint/compliance/checker.py

@@ -157,6 +157,21 @@ def check_config_lock(config: dict[str, Any]) -> list[Check]:
         detail += f" target_concurrency={target_concurrency}"
     checks.append(Check("single_stream", single_stream, detail))
 
+    # Lock the pre-response reset retry: it is a reproducibility knob that lets
+    # the client silently re-issue after a SUT-side reset, so its value must be
+    # recorded (and within the schema bound) rather than passing for any value.
+    # Disclosure that a recovered reset didn't mask a dropped turn is a manual
+    # attestation added in check_submission (transport_retries_total lives in the
+    # run's result summary, not in this config).
+    retries = _get(config, "settings", "client", "transport_max_retries", default=0)
+    checks.append(
+        Check(
+            "transport_max_retries_locked",
+            isinstance(retries, int) and 0 <= retries <= 5,
+            f"transport_max_retries={retries}",
+        )
+    )
+
     return checks
 
 
@@ -309,6 +324,22 @@ def check_submission(
             report.add("config_valid", False, f"{config_path} failed to load: {e}")
         else:
             report.checks.extend(check_config_lock(config))
+            retries = _get(
+                config, "settings", "client", "transport_max_retries", default=0
+            )
+            if retries:
+                # A recovered pre-response reset produces a clean success, so
+                # the no_dropped_turns gate can't see it from scores.json alone.
+                # Require the submitter to confirm the disclosed retry value and
+                # that transport_retries_total (in the run's result summary) did
+                # not mask a dropped-and-re-issued turn.
+                report.notes.append(
+                    f"transport_max_retries={retries}: pre-response connection "
+                    "resets are retried. Confirm this matches the disclosed "
+                    "ruleset value and that no retry-recovered reset (see "
+                    "transport_retries_total in the run's result summary) masked "
+                    "a dropped turn."
+                )
 
     results_path = report_dir / "results.json"
     scores_path = report_dir / "scores.json"

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, max 5)
     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, max 5)
     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, max 5)
     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/core/record.py

@@ -145,6 +145,9 @@ class SampleEventType(EventType):
     COMPLETE = "complete"
     RECV_FIRST = "recv_first"
     RECV_NON_FIRST = "recv_non_first"
+    # Emitted once per pre-response connection-reset re-issue (opt-in
+    # transport_max_retries). Diagnostic only — not a tracked sample event.
+    TRANSPORT_RETRY = "transport_retry"
 
 
 class EventRecord(msgspec.Struct, kw_only=True, frozen=True, gc=False):  # type: ignore[call-arg]

src/inference_endpoint/core/types.py

@@ -402,6 +402,12 @@ class QueryResult(
         error: Structured error if query failed (None if successful).
         completed_at: High-resolution timestamp (nanoseconds, monotonic clock).
                       Auto-set in __post_init__ to prevent tampering.
+        transport_retries: Number of pre-response connection-reset re-issues the
+                      worker performed for this sample (0 = none). Non-zero marks
+                      a sample whose timed window folds in dead-socket wait +
+                      reconnect + a second server pass, so the report can
+                      segregate or exclude it from the server latency
+                      distribution. Appended last for array_like backward compat.
 
     Note:
         The completed_at field is intentionally set internally to prevent
@@ -422,6 +428,7 @@ class QueryResult(
     metadata: dict[str, Any] = msgspec.field(default_factory=dict)
     error: ErrorData | None = None
     completed_at: int | msgspec.UnsetType = msgspec.UNSET
+    transport_retries: int = 0
 
     def __post_init__(self):
         """Set completion timestamp automatically.

src/inference_endpoint/endpoint_client/config.py

@@ -170,6 +170,26 @@ 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`.
+    # Bounded above so a genuinely-down server can't turn fail-fast into a long
+    # reconnect/re-write spin that masks a hard SUT failure and multiplies load.
+    transport_max_retries: int = Field(
+        0,
+        ge=0,
+        le=5,
+        description="Retries on a pre-response connection reset (0=disabled, max 5)",
+    )
+
     # 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/http.py

@@ -131,6 +131,7 @@ class HttpResponseProtocol(asyncio.Protocol):
         "_message_complete",
         "_connection_lost",
         "_exc",
+        "_bytes_received",
     )
 
     def __init__(self, loop: asyncio.AbstractEventLoop):
@@ -158,6 +159,10 @@ def __init__(self, loop: asyncio.AbstractEventLoop):
         self._message_complete: bool = False
         self._connection_lost: bool = False
         self._exc: Exception | None = None
+        # True once any response byte has been received for the current
+        # request. Distinguishes a true zero-byte close (safe to retry) from a
+        # partial-response close (server already processed the request).
+        self._bytes_received: bool = False
 
     def reset(self) -> None:
         """Reset protocol state for connection reuse."""
@@ -175,6 +180,7 @@ def reset(self) -> None:
         self._headers_complete = False
         self._message_complete = False
         self._exc = None
+        self._bytes_received = False
         # NOTE: Don't reset _connection_lost - that's transport state
 
     def _signal_stream_end(self) -> None:
@@ -196,6 +202,8 @@ def connection_made(self, transport: asyncio.Transport) -> None:  # type: ignore
         self._parser = httptools.HttpResponseParser(self)
 
     def data_received(self, data: bytes) -> None:
+        if data:
+            self._bytes_received = True
         # Lazy parser creation for better reset() performance
         if self._parser is None:
             self._parser = httptools.HttpResponseParser(self)
@@ -316,6 +324,17 @@ def should_close(self) -> bool:
         """Whether connection should be closed after this response."""
         return self._should_close or self._connection_lost or self._exc is not None
 
+    @property
+    def bytes_received(self) -> bool:
+        """Whether any response byte arrived for the current request.
+
+        False means a true zero-byte close (the server never wrote a byte);
+        True means partial or complete response bytes were parsed. Used to gate
+        the pre-response reset retry so a request the server may already have
+        processed is not re-issued.
+        """
+        return self._bytes_received
+
     def write(self, data: bytes) -> None:
         """Write data to transport."""
         if self._transport:
@@ -588,8 +607,11 @@ def release(self, conn: PooledConnection) -> None:
         if not conn.in_use:
             return
 
-        # Must close if: dead, server requested close, or error occurred
+        # Must close if: dead, server requested close, or error occurred.
+        # Clear in_use first so a second release() short-circuits above
+        # instead of re-entering this branch and firing a spurious waiter wake.
         if not conn.is_alive() or conn.protocol.should_close:
+            conn.in_use = False
             self._close_connection(conn)
             self._notify_waiter()
             return
@@ -798,9 +820,11 @@ class InFlightRequest:
         http_bytes: Serialized HTTP request for socket.write().
         is_streaming: Whether this is a streaming (SSE) request or not.
         connection: PooledConnection assigned to this request (set once request is fired).
+        transport_retries: Count of pre-response reset re-issues for this request.
     """
 
     query_id: str
     http_bytes: bytes
     is_streaming: bool
     connection: PooledConnection = field(default=None, repr=False)  # type: ignore[assignment]
+    transport_retries: int = 0

src/inference_endpoint/endpoint_client/worker.py

@@ -28,6 +28,8 @@
 from typing import Any
 from urllib.parse import urlparse
 
+import msgspec
+
 from inference_endpoint.async_utils.transport import (
     ReceiverTransport,
     SenderTransport,
@@ -378,14 +380,64 @@ 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 only on a true zero-byte close.
+
+        A connection-level failure here can mean the server closed the socket
+        before writing any response byte (the idle keep-alive race: an idle
+        ``-np 1`` connection dropped just as the client writes on it). Only a
+        true zero-byte close is retried — re-issuing on a fresh connection is
+        safe for *client* output because the accumulator has produced nothing.
+
+        If any response byte already arrived (partial status-line/header bytes
+        before the close), the server may have received and processed the
+        request, so a re-issue would duplicate server-side generation work; that
+        case surfaces the error instead of retrying. Each recovery is logged and
+        counted on ``req.transport_retries`` so the successful sample carries a
+        marker downstream. Returns the HTTP status code, or None if the reset is
+        not retryable or retries are exhausted (an error response was emitted).
+        """
+        while True:
+            conn = req.connection
+            try:
+                status_code, _ = await conn.protocol.read_headers()
+                return status_code
+            except ConnectionError as e:
+                partial = conn.protocol.bytes_received
+                self._pool.release(conn)
+                if (
+                    partial
+                    or req.transport_retries >= self.http_config.transport_max_retries
+                ):
+                    await self._handle_error(req.query_id, e)
+                    return None
+                req.transport_retries += 1
+                logger.warning(
+                    "Request %s: retrying after pre-response connection reset "
+                    "(attempt %d/%d): %s",
+                    req.query_id,
+                    req.transport_retries,
+                    self.http_config.transport_max_retries,
+                    e,
+                )
+                # Attach the fresh connection to the request BEFORE writing, so
+                # a write failure still leaves req.connection pointing at it for
+                # the caller's finally-block release (no leaked connection).
+                new_conn = await self._pool.acquire()
+                req.connection = new_conn
+                new_conn.protocol.write(req.http_bytes)
+
     @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 +458,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()
@@ -432,7 +485,7 @@ async def _handle_streaming_body(self, req: InFlightRequest) -> None:
         self._pool.release(conn)
 
         # Send final complete back to main rank
-        self._responses.send(accumulator.get_final_output())
+        self._responses.send(self._tag_retries(accumulator.get_final_output(), req))
 
     @profile
     async def _handle_non_streaming_body(self, req: InFlightRequest) -> None:
@@ -450,7 +503,20 @@ async def _handle_non_streaming_body(self, req: InFlightRequest) -> None:
         result = self._adapter.decode_response(response_bytes, query_id)
 
         # Send result back to main rank
-        self._responses.send(result)
+        self._responses.send(self._tag_retries(result, req))
+
+    @staticmethod
+    def _tag_retries(result: QueryResult, req: InFlightRequest) -> QueryResult:
+        """Stamp a successful result with this request's retry count.
+
+        No-op when the request was not retried (the common path), so the hot
+        path pays nothing unless a pre-response reset was actually recovered.
+        """
+        if req.transport_retries:
+            msgspec.structs.force_setattr(
+                result, "transport_retries", req.transport_retries
+            )
+        return result
 
     async def _handle_error(self, query_id: str, error: Exception | str) -> None:
         """Send error response for a query."""

src/inference_endpoint/load_generator/session.py

@@ -600,6 +600,21 @@ def _handle_response(self, resp: QueryResult | StreamChunk) -> None:
                         data=resp.error,
                     )
                 )
+            # Surface pre-response connection-reset recoveries as a diagnostic
+            # counter (one event per re-issue). Emitted for every completed
+            # sample, warmup included, since a reset is a SUT-health signal
+            # regardless of the tracking window.
+            if resp.transport_retries:
+                retry_record = EventRecord(
+                    event_type=SampleEventType.TRANSPORT_RETRY,
+                    timestamp_ns=time.monotonic_ns(),
+                    sample_uuid=query_id,
+                    conversation_id=conv_id_str,
+                    turn=turn_num,
+                )
+                for _ in range(resp.transport_retries):
+                    self._publisher.publish(retry_record)
+
             if self._current_phase_type != PhaseType.WARMUP:
                 self._publisher.publish(
                     EventRecord(