[BUG CLIENT]: Unhandled transient network errors crash long-running BatchClient.wait() polling loops

[andriiorap]

Python -VV

Python 3.11.8 (default, Feb 12 2024, 14:50:05) [Clang 15.0.0 (clang-1500.1.0.2.5)]
(Applies to all Python versions: 3.9, 3.10, 3.11, 3.12, 3.13, 3.14)

Pip Freeze

mistralai==2.10.0
httpx>=0.27.0
pydantic>=2.0.0

Reproduction Steps

1. Launch a long-running batch job using BatchClient.wait() or BatchClient.run() / BatchClient.create_and_wait() with a realistic multi-hour duration (e.g., default timeout_hours=24, poll_interval_seconds=30.0):

   import asyncio
   from mistralai import Mistral
   from mistralai.extra.batch import BatchClient
   
   async def main():
       client = Mistral()
       batch = BatchClient(client)
   
       # Simulate or start a long batch job
       handle = await batch.create(...)
       # Wait for completion (default up to 24 hours)
       result = await batch.wait(handle, poll_interval_seconds=30.0, timeout_hours=24)
   
   asyncio.run(main())

2. During the 24-hour polling period (~2,880 status check HTTP requests), simulate a brief transient network glitch (e.g. 502/503/504 Bad Gateway during edge/gateway deployments, momentary DNS resolution timeout, or httpx.ConnectError / httpx.ReadTimeout).
3. Notice that batch.wait() immediately raises the unhandled network exception and terminates execution, crashing the entire caller application or workflow pipeline.

Expected Behavior

Batch jobs are explicitly designed for asynchronous, long-duration batch processing (up to 24 hours). Over thousands of polling HTTP requests spanning hours, transient network blips, gateway 502/503s, or momentary rate-limit (429) spikes are statistically inevitable.

BatchClient.wait() should handle transient transport/server errors gracefully during polling:

• Log a warning when a transient error occurs during a poll cycle.
• Continue polling with backoff until the deadline is reached or until a threshold of consecutive failures (e.g. 5-10 consecutive failed attempts) indicates an unrecoverable outage.
• A single dropped packet or 502 gateway error during hour 18 should not crash the entire job wait and orphan the batch job.

Actual Behavior

In src/mistralai/extra/batch/client.py (lines 660–668):

    async def wait(
        self,
        handle: Union["BatchJobHandle[RespBodyT]", str],
        *,
        poll_interval_seconds: float = 30.0,
        timeout_hours: int = 24,
        deadline: Optional[float] = None,
        http_headers: Optional[Mapping[str, str]] = None,
    ) -> "BatchJobHandle[RespBodyT]":
        job_id = self._job_id(handle)
        if deadline is None:
            deadline = time.monotonic() + timeout_hours * 3600
        while True:
            current = await self.get(job_id, http_headers=http_headers)  # <-- UNGUARDED
            if is_terminal(current.status):
                return current
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError(f"Batch job {job_id} timed out while polling")
            await asyncio.sleep(min(poll_interval_seconds, remaining))

The call to self.get(job_id, http_headers=http_headers) is completely unshielded by any try...except. If self.get() raises httpx.TransportError, httpx.TimeoutException, or a transient 5xx errors.SDKError, the exception immediately propagates out of wait(), crashing downstream callers such as create_and_wait(), run(), and user pipelines.

Additional Context

• The codebase already recognizes that network errors should not mask outcomes in other places. For example, in BatchClient._delete_quietly (lines 936–941):

  """Best-effort file delete for run()'s cleanup -- a failed delete must never
  mask the real result (or the real exception on the abort path)."""

• For multi-hour batch workloads, having zero error tolerance in the status polling loop creates severe fragility in automated ETL and ML production pipelines.

Suggested Solutions

A production-grade retry implementation requires three specific refinements:

1. Filter HTTP Status Codes (Do Not Retry Indiscriminately):
   • Retryable (Transient Errors): Status codes 408, 429, 500, 502, 503, 504, along with httpx.TransportError (network drops, connection resets, DNS failures) and httpx.TimeoutException.
   • Fail Immediately (Permanent Errors): 400, 401, 403, 404, 422. If the API key is revoked or the job ID does not exist, repeatedly polling wastes time, burns retries, and delays error discovery.
2. Respect the Remaining Deadline During Error Backoff:
   • The sleep duration must always be clamped to remaining = deadline - time.monotonic(). If remaining <= 0, immediately raise TimeoutError.
3. Exponential Backoff with Jitter & Retry-After Header Respect:
   • For consecutive transient errors, back off exponentially with jitter to prevent thundering herd spikes against a recovering gateway.
   • On 429 / 503, inspect and honor the server's Retry-After header when present.

Refined Implementation:

import asyncio
import random
import time
from typing import Optional, Union, Mapping
import httpx
from mistralai.client.errors import MistralError

RETRYABLE_STATUS_CODES = {408, 429, 500, 502, 503, 504}


def _is_transient_error(exc: Exception) -> bool:
    """Identify transient transport or gateway errors that are safe to retry."""
    if isinstance(exc, (httpx.TransportError, httpx.TimeoutException)):
        return True
    if isinstance(exc, MistralError):
        return exc.status_code in RETRYABLE_STATUS_CODES
    if isinstance(exc, httpx.HTTPStatusError):
        return exc.response.status_code in RETRYABLE_STATUS_CODES
    return False


def _get_retry_after(exc: Exception) -> Optional[float]:
    """Extract server-instructed Retry-After header in seconds, if available."""
    if isinstance(exc, MistralError) and exc.headers:
        val = exc.headers.get("retry-after")
        if val is not None:
            try:
                return max(0.0, float(val))
            except ValueError:
                pass
    return None


async def wait(
    self,
    handle: Union["BatchJobHandle[RespBodyT]", str],
    *,
    poll_interval_seconds: float = 30.0,
    timeout_hours: int = 24,
    deadline: Optional[float] = None,
    http_headers: Optional[Mapping[str, str]] = None,
    max_consecutive_errors: int = 10,
) -> "BatchJobHandle[RespBodyT]":
    """Poll until the job reaches a terminal state, returning the terminal handle.
    Pass ``deadline`` (a ``time.monotonic()`` value) to enforce a wall-clock cap
    shared across several calls; otherwise it is derived from ``timeout_hours``."""
    job_id = self._job_id(handle)
    if deadline is None:
        deadline = time.monotonic() + timeout_hours * 3600

    consecutive_errors = 0

    while True:
        try:
            current = await self.get(job_id, http_headers=http_headers)
            consecutive_errors = 0
            if is_terminal(current.status):
                return current
            sleep_duration = poll_interval_seconds
        except Exception as exc:
            # 1. Filter HTTP status codes: fail fast on 400, 401, 403, 404, 422
            if not _is_transient_error(exc):
                raise

            consecutive_errors += 1
            if consecutive_errors >= max_consecutive_errors:
                raise

            # 3. Honor Retry-After header or apply exponential backoff with jitter
            retry_after = _get_retry_after(exc)
            if retry_after is not None:
                sleep_duration = retry_after
            else:
                base_backoff = min(poll_interval_seconds, 2.0 ** (consecutive_errors - 1))
                jitter = random.uniform(0.8, 1.2)
                sleep_duration = base_backoff * jitter

        # 2. Respect remaining deadline
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise TimeoutError(f"Batch job {job_id} timed out while polling")

        await asyncio.sleep(min(sleep_duration, remaining))