fix(python): retry transient job polling errors

Fixes #4773.

Python job polling now retries transient connection, send, and read failures with capped exponential backoff under the caller's original timeout. If the deadline expires while disconnected, `wait()` raises `ConnectionError` with the job ID, reconnect guidance, and the original cause; after recovery, ordinary timeouts remain `TimeoutError`. Terminal job errors, cancellation, and malformed responses continue to surface immediately.

Validation: exact-head CI passed across Rust all-features, Python on Linux, macOS, and Windows, NodeJS, deny, and packaging checks. The hosted macOS wheel suite passed 3049 tests (162 skipped, 70 deselected), including sync and async recovery, timeout reuse, cancellation, terminal errors, and error preservation.

python/Cargo.toml

@@ -36,6 +36,7 @@ pyo3-async-runtimes = { version = "0.29", features = [
     "attributes",
     "tokio-runtime",
 ] }
+reqwest.workspace = true
 pin-project.workspace = true
 futures.workspace = true
 serde.workspace = true

python/python/lancedb/job.py

@@ -14,6 +14,7 @@ from lancedb.background_loop import LOOP
 
 from . import _lancedb
 from ._lancedb import JobDescription, JobFailureInfo, JobInfo
+from .remote.errors import HttpError, RetryError
 
 T = TypeVar("T")
 
@@ -95,6 +96,12 @@ class AsyncJob(Generic[T]):
 
         Raises `JobFailedError` if the operation failed, `JobCancelledError`
         if it was cancelled, and `TimeoutError` if `timeout` elapses first.
+        Transient connection and read errors are retried with capped exponential
+        backoff within the same timeout (indefinitely when no timeout is given).
+        If the timeout expires while disconnected, raises `ConnectionError`
+        instead, with the job ID and instructions to resume waiting. The job
+        may still be running; timing out does not cancel it.
+
         The timeout error identifies the job and explains that it continues on
         the server. Registry jobs include the ``db.open_job(id).wait()``
         reattach call; SQL query handles must be kept because their IDs cannot
@@ -102,17 +109,55 @@ class AsyncJob(Generic[T]):
         """
         if self._inner is None:
             return cast(T, None)
-        if timeout is None:
-            result = await self._inner.wait()
-        else:
-            try:
+        last_error = None
+
+        async def wait_with_retries():
+            nonlocal last_error
+            delay = 0.25
+            while True:
+                try:
+                    if last_error is not None:
+                        # Confirm recovery before restarting the backend poller,
+                        # so a later timeout isn't blamed on an earlier outage.
+                        await self._inner.status()
+                        last_error = None
+                        delay = 0.25
+                    return await self._inner.wait()
+                except (HttpError, RetryError) as error:
+                    cause = error.__cause__ if isinstance(error, RetryError) else error
+                    if not isinstance(cause, HttpError) or not cause.is_transport_error:
+                        raise
+                    last_error = error
+                await asyncio.sleep(delay)
+                delay = min(delay * 2, 30.0)
+
+        try:
+            if timeout is None:
+                result = await wait_with_retries()
+            else:
                 result = await asyncio.wait_for(
-                    self._inner.wait(), timeout.total_seconds()
+                    wait_with_retries(), timeout.total_seconds()
                 )
-            except asyncio.TimeoutError as error:
+        except asyncio.TimeoutError as error:
+            if last_error is None:
+                if timeout is None:
+                    raise
                 raise TimeoutError(
                     _timeout_message(self.id, timeout, self.job_type)
                 ) from error
+            identity = f"job {self.id!r}" if self.id is not None else "the operation"
+            if self.id is not None and self.job_type != "query":
+                reconnect = (
+                    f"Use db.open_job({self.id!r}) to reconnect, then call wait() "
+                    "on the returned job."
+                )
+            else:
+                reconnect = "Call wait() on this job handle to reconnect."
+            raise ConnectionError(
+                f"Lost connection while waiting and the timeout expired; "
+                f"{identity} may still be running. "
+                f"{reconnect}"
+            ) from last_error
         if self._result_decoder is not None:
             return self._result_decoder(result)
         return cast(T, result)
@@ -303,10 +348,16 @@ class Job(Generic[T]):
 
         Raises `JobFailedError` if the operation failed, `JobCancelledError`
         if it was cancelled, and `TimeoutError` if `timeout` elapses first.
+        Transient connection and read errors are retried within that timeout;
+        if it expires while disconnected, raises `ConnectionError` with
+        instructions to reconnect. Timing out does not cancel the job.
+
         The timeout error identifies the job and explains that it continues on
         the server. Registry jobs include the ``db.open_job(id).wait()``
         reattach call; SQL query handles must be kept because their IDs cannot
         be reopened through the registry jobs API.
+
+        See :meth:`AsyncJob.wait`.
         """
         if self._inner is None:
             return cast(T, None)

python/python/lancedb/remote/errors.py

@@ -54,6 +54,9 @@ class HttpError(LanceDBClientError):
         The name of the resource reported by the server, if available.
     namespace_path: tuple[str, ...]
         The parent namespace path of the resource, if available.
+    is_transport_error: bool
+        Whether the request failed while connecting, sending, or reading, rather
+        than because of an HTTP status or an invalid response payload.
     """
 
     def __new__(
@@ -65,6 +68,7 @@ class HttpError(LanceDBClientError):
         resource_kind: Optional[str] = None,
         resource_name: Optional[str] = None,
         namespace_path: Optional[tuple[str, ...]] = None,
+        is_transport_error: bool = False,
     ):
         if cls is HttpError:
             mapped_cls = _http_error_class(error_code, resource_kind)
@@ -81,8 +85,10 @@ class HttpError(LanceDBClientError):
         resource_kind: Optional[str] = None,
         resource_name: Optional[str] = None,
         namespace_path: Optional[tuple[str, ...]] = None,
+        is_transport_error: bool = False,
     ):
         super().__init__(message, request_id, status_code)
+        self.is_transport_error = is_transport_error
         self.error_code = error_code
         self.resource_kind = resource_kind
         self.resource_name = resource_name
@@ -101,6 +107,7 @@ class HttpError(LanceDBClientError):
                 self.resource_kind,
                 self.resource_name,
                 self.namespace_path,
+                self.is_transport_error,
             ),
         )
 

python/python/tests/test_errors.py

@@ -61,6 +61,15 @@ def test_pickle_http_error():
     assert str(restored) == "not found"
     assert restored.request_id == "req-789"
     assert restored.status_code == 404
+    assert restored.is_transport_error is False
+
+
+def test_pickle_http_transport_error():
+    err = HttpError("connection reset", "req-789", is_transport_error=True)
+    restored = pickle.loads(pickle.dumps(err))
+    assert isinstance(restored, HttpError)
+    assert restored.is_transport_error is True
+    assert restored.request_id == "req-789"
 
 
 @pytest.mark.parametrize(
@@ -93,6 +102,7 @@ def test_http_error_structured_resource_error(
     assert err.namespace_path == ("parent", "child")
     assert err.request_id == "req-structured"
     assert err.status_code == status_code
+    assert err.is_transport_error is False
 
 
 @pytest.mark.parametrize(
@@ -131,7 +141,10 @@ def test_http_error_without_known_structured_fields_stays_generic(
         (AlreadyExistsError, "already_exists", "namespace"),
     ],
 )
-def test_pickle_structured_http_error(exception_type, error_code, resource_kind):
+@pytest.mark.parametrize("is_transport_error", [False, True])
+def test_pickle_structured_http_error(
+    exception_type, error_code, resource_kind, is_transport_error
+):
     err = HttpError(
         "resource request failed",
         "req-pickle",
@@ -140,6 +153,7 @@ def test_pickle_structured_http_error(exception_type, error_code, resource_kind)
         resource_kind,
         "resource-name",
         ["parent"],
+        is_transport_error=is_transport_error,
     )
     assert type(err) is exception_type
 
@@ -152,6 +166,7 @@ def test_pickle_structured_http_error(exception_type, error_code, resource_kind)
     assert restored.namespace_path == ("parent",)
     assert restored.request_id == "req-pickle"
     assert restored.status_code == err.status_code
+    assert restored.is_transport_error is is_transport_error
 
 
 def test_pickle_table_not_found_error():

python/python/tests/test_job.py

@@ -3,10 +3,223 @@
 
 import asyncio
 from datetime import timedelta
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
+from lancedb.exceptions import JobCancelledError, JobFailedError
 from lancedb.job import AsyncJob, Job
+from lancedb.remote.errors import HttpError, RetryError
+
+
+@pytest.fixture(params=["async", "sync"])
+def wait_job(request):
+    async def wait(inner, timeout=None, result_decoder=None):
+        job = AsyncJob(inner, result_decoder)
+        if request.param == "sync":
+            return await asyncio.to_thread(Job(job).wait, timeout)
+        return await job.wait(timeout)
+
+    return wait
+
+
+def transport_error(kind):
+    if kind == "http":
+        return HttpError("error sending request", "req-1", is_transport_error=True)
+    error = RetryError("retries exhausted", "req-1", 0, 3, 0, 3, 3, 3, None)
+    error.__cause__ = transport_error("http")
+    return error
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", ["http", "retry"])
+async def test_wait_recovers_from_transport_errors(wait_job, kind, monkeypatch):
+    inner = MagicMock(id="job-1", job_type=None)
+    inner.wait = AsyncMock(side_effect=[transport_error(kind), '{"rows": 7}'])
+    inner.status = AsyncMock(side_effect=[transport_error(kind), "running"])
+    sleep = AsyncMock()
+    monkeypatch.setattr("lancedb.job.asyncio.sleep", sleep)
+
+    result = await wait_job(
+        inner, timedelta(hours=24), result_decoder=lambda value: ("decoded", value)
+    )
+
+    assert result == ("decoded", '{"rows": 7}')
+    assert inner.wait.await_count == 2
+    assert inner.status.await_count == 2
+    assert [call.args[0] for call in sleep.await_args_list] == [0.25, 0.5]
+    inner.cancel.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_wait_disconnection_timeout_explains_how_to_reconnect(wait_job):
+    error = transport_error("http")
+    inner = MagicMock(id="job-1", job_type=None)
+    inner.wait = AsyncMock(side_effect=error)
+    inner.status = AsyncMock(side_effect=error)
+
+    with pytest.raises(
+        ConnectionError, match="job 'job-1' may still be running"
+    ) as exc:
+        await wait_job(inner, timedelta(milliseconds=10))
+
+    assert "db.open_job('job-1')" in str(exc.value)
+    assert "call wait()" in str(exc.value)
+    assert exc.value.__cause__ is error
+    inner.cancel.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "error",
+    [
+        HttpError("unauthorized", "req-1", 401),
+        HttpError("not found", "req-1", 404),
+        HttpError(
+            "secret not found",
+            "req-1",
+            404,
+            error_code="not_found",
+            resource_kind="secret",
+        ),
+        HttpError("failed to parse job description", "req-1"),
+        RetryError("service unavailable", "req-1", 3, 0, 0, 3, 3, 3, 503),
+        JobFailedError("worker failed"),
+        JobCancelledError("job cancelled"),
+    ],
+)
+async def test_wait_propagates_permanent_and_terminal_errors(wait_job, error):
+    inner = MagicMock(id="job-1")
+    inner.wait = AsyncMock(side_effect=error)
+
+    with pytest.raises(type(error)) as exc:
+        await wait_job(inner, timedelta(seconds=1))
+
+    assert exc.value is error
+    assert inner.wait.await_count == 1
+    inner.status.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_wait_running_job_timeout_is_not_a_connection_error(wait_job):
+    async def running():
+        await asyncio.Event().wait()
+
+    inner = MagicMock(id="job-1")
+    inner.wait = AsyncMock(side_effect=running)
+    with pytest.raises(TimeoutError):
+        await wait_job(inner, timedelta(milliseconds=10))
+    inner.cancel.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_wait_completed_job(wait_job):
+    assert await wait_job(None) is None
+
+
+@pytest.mark.asyncio
+async def test_wait_retries_without_a_timeout_and_caps_backoff(wait_job, monkeypatch):
+    inner = MagicMock(id="job-1", job_type=None)
+    inner.wait = AsyncMock(side_effect=[transport_error("http"), None])
+    inner.status = AsyncMock(side_effect=[transport_error("retry")] * 9 + ["running"])
+    sleep = AsyncMock()
+    monkeypatch.setattr("lancedb.job.asyncio.sleep", sleep)
+
+    assert await wait_job(inner) is None
+    assert [call.args[0] for call in sleep.await_args_list] == [
+        0.25,
+        0.5,
+        1,
+        2,
+        4,
+        8,
+        16,
+        30,
+        30,
+        30,
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("job_type", [None, "query"])
+async def test_wait_timeout_after_recovery_is_not_a_connection_error(
+    wait_job, monkeypatch, job_type
+):
+    async def running():
+        await asyncio.Event().wait()
+
+    inner = MagicMock(id="job-1", job_type=job_type)
+    inner.wait = AsyncMock()
+
+    async def wait():
+        if inner.wait.await_count == 1:
+            raise transport_error("http")
+        return await running()
+
+    inner.wait.side_effect = wait
+    inner.status = AsyncMock(return_value="running")
+    monkeypatch.setattr("lancedb.job.asyncio.sleep", AsyncMock())
+
+    with pytest.raises(TimeoutError) as exc:
+        await wait_job(inner, timedelta(milliseconds=10))
+    message = str(exc.value)
+    assert "still running after 0.01 seconds" in message
+    assert "continues on the server" in message
+    if job_type == "query":
+        assert "keep this handle to continue waiting" in message
+        assert "db.open_job" not in message
+    else:
+        assert "db.open_job('job-1').wait()" in message
+    assert inner.status.await_count == 1
+    assert inner.wait.await_count == 2
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "error", [JobFailedError("failed"), JobCancelledError("cancel")]
+)
+async def test_wait_reports_terminal_error_after_recovery(wait_job, monkeypatch, error):
+    inner = MagicMock(id="job-1", job_type=None)
+    inner.wait = AsyncMock(side_effect=[transport_error("retry"), error])
+    inner.status = AsyncMock(return_value="running")
+    monkeypatch.setattr("lancedb.job.asyncio.sleep", AsyncMock())
+
+    with pytest.raises(type(error)) as exc:
+        await wait_job(inner, timedelta(seconds=1))
+    assert exc.value is error
+    assert inner.wait.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_wait_cancellation_interrupts_backoff(monkeypatch):
+    inner = MagicMock(id="job-1", job_type=None)
+    inner.wait = AsyncMock(side_effect=transport_error("http"))
+    sleeping = asyncio.Event()
+
+    async def sleep(delay):
+        sleeping.set()
+        await asyncio.Event().wait()
+
+    monkeypatch.setattr("lancedb.job.asyncio.sleep", sleep)
+    task = asyncio.create_task(AsyncJob(inner).wait())
+    await asyncio.wait_for(sleeping.wait(), 1)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    inner.cancel.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_query_wait_disconnection_does_not_suggest_open_job(wait_job):
+    inner = MagicMock(id="query-1", job_type="query")
+    inner.wait = AsyncMock(side_effect=transport_error("http"))
+
+    with pytest.raises(ConnectionError) as exc:
+        await wait_job(inner, timedelta(milliseconds=10))
+
+    assert "query-1" in str(exc.value)
+    assert "this job handle" in str(exc.value)
+    assert "open_job" not in str(exc.value)
 
 
 class _PendingJob:

python/python/tests/test_remote_db.py

@@ -12,6 +12,7 @@ import json
 import multiprocessing as mp
 import pickle
 import re
+import socket
 import subprocess
 import sys
 import textwrap
@@ -162,6 +163,7 @@ def test_remote_open_missing_table(cached, namespace_path, missing_table_handler
     assert str(error) == "Table 'missing' was not found"
     assert error.status_code == 404
     assert error.request_id == request_ids[-1]
+    assert error.is_transport_error is False
     assert error.request_id
     assert isinstance(error.__cause__, HttpError)
     assert error.__cause__.request_id == error.request_id
@@ -337,6 +339,7 @@ def test_remote_create_namespace_typed_already_exists():
     assert error.resource_name == "child"
     assert error.namespace_path == ("parent",)
     assert error.status_code == 409
+    assert error.is_transport_error is False
     assert len(requests) == 1
 
 
@@ -1151,6 +1154,7 @@ async def test_http_error():
 
         assert exc_info.value.request_id == request_id_holder["request_id"]
         assert "Internal Server Error" in str(exc_info.value)
+        assert exc_info.value.is_transport_error is False
 
 
 @pytest.mark.asyncio
@@ -1175,6 +1179,24 @@ async def test_retry_error():
         assert "Try again later" in str(cause)
         assert cause.request_id == request_id_holder["request_id"]
         assert cause.status_code == 429
+        assert cause.is_transport_error is False
+
+
+def test_connect_retry_error_preserves_transport_classification():
+    # Reserve a port without listening, so connection attempts are refused.
+    with socket.socket() as port_guard:
+        port_guard.bind(("127.0.0.1", 0))
+        port = port_guard.getsockname()[1]
+        db = lancedb.connect(
+            "db://dev",
+            api_key="fake",
+            host_override=f"http://127.0.0.1:{port}",
+            client_config={"retry_config": {"connect_retries": 0}},
+        )
+        with pytest.raises(RetryError) as exc:
+            list(db.list_tables())
+        assert isinstance(exc.value.__cause__, HttpError)
+        assert exc.value.__cause__.is_transport_error is True
 
 
 def test_table_unimplemented_functions():
@@ -1931,6 +1953,114 @@ def test_remote_function_errors_lists_the_rows_a_refresh_skipped():
     assert everything.truncated is True
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [False, True])
+@pytest.mark.parametrize("failure", ["send", "read"])
+async def test_remote_job_wait_recovers_from_disconnection(sync, failure):
+    polls = 0
+
+    def handler(request):
+        nonlocal polls
+        assert request.path == "/v1/jobs/describe"
+        body = request.rfile.read(int(request.headers["Content-Length"]))
+        assert json.loads(body)["job_id"] == "job-1"
+        polls += 1
+        if polls in (2, 3):
+            # Fail both the initial wait and its first recovery probe, after
+            # open_job has successfully fetched the still-running job.
+            if failure == "read":
+                request.send_response(200)
+                request.send_header("Content-Length", "1000")
+                request.end_headers()
+                request.wfile.write(b'{"job_id":')
+                request.wfile.flush()
+            request.connection.shutdown(socket.SHUT_RDWR)
+            request.close_connection = True
+            return
+        request.send_response(200)
+        request.send_header("Content-Type", "application/json")
+        request.end_headers()
+        state = "DONE" if polls >= 5 else "IN_PROGRESS"
+        request.wfile.write(json.dumps(dict(job_id="job-1", job_state=state)).encode())
+
+    config = {"retry_config": {"retries": 0, "connect_retries": 0, "read_retries": 0}}
+    if sync:
+        with mock_lancedb_connection(handler, **config) as db:
+            job = db.open_job("job-1")
+            assert job.wait(timeout=timedelta(seconds=5)) is None
+            assert job.state == "finished"
+    else:
+        async with mock_lancedb_connection_async(handler, **config) as db:
+            job = await db.open_job("job-1")
+            assert await job.wait(timeout=timedelta(seconds=5)) is None
+            assert job.state == "finished"
+    assert polls == 5
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [False, True])
+async def test_remote_job_wait_disconnection_timeout_can_be_resumed(sync):
+    polls = 0
+    finished = False
+
+    def handler(request):
+        nonlocal polls
+        assert request.path == "/v1/jobs/describe"
+        request.rfile.read(int(request.headers["Content-Length"]))
+        polls += 1
+        if polls > 1 and not finished:
+            request.connection.shutdown(socket.SHUT_RDWR)
+            request.close_connection = True
+            return
+        request.send_response(200)
+        request.end_headers()
+        state = "DONE" if finished else "IN_PROGRESS"
+        request.wfile.write(json.dumps(dict(job_id="job-1", job_state=state)).encode())
+
+    if sync:
+        with mock_lancedb_connection(handler) as db:
+            job = db.open_job("job-1")
+            with pytest.raises(ConnectionError) as exc:
+                job.wait(timeout=timedelta(seconds=0.5))
+            finished = True
+            assert db.open_job("job-1").wait() is None
+    else:
+        async with mock_lancedb_connection_async(handler) as db:
+            job = await db.open_job("job-1")
+            with pytest.raises(ConnectionError) as exc:
+                await job.wait(timeout=timedelta(seconds=0.5))
+            finished = True
+            resumed = await db.open_job("job-1")
+            assert await resumed.wait() is None
+    assert "job 'job-1' may still be running" in str(exc.value)
+    assert "db.open_job('job-1')" in str(exc.value)
+    assert isinstance(exc.value.__cause__, HttpError)
+    assert exc.value.__cause__.is_transport_error is True
+
+
+@pytest.mark.parametrize("payload", [b"{", b"{}"])
+def test_remote_job_wait_does_not_retry_malformed_description(payload):
+    polls = 0
+
+    def handler(request):
+        nonlocal polls
+        assert request.path == "/v1/jobs/describe"
+        request.rfile.read(int(request.headers["Content-Length"]))
+        polls += 1
+        request.send_response(200)
+        request.end_headers()
+        request.wfile.write(
+            b'{"job_id":"job-1","job_state":"IN_PROGRESS"}' if polls == 1 else payload
+        )
+
+    with mock_lancedb_connection(handler) as db:
+        job = db.open_job("job-1")
+        with pytest.raises(HttpError, match="failed to parse job description") as exc:
+            job.wait(timeout=timedelta(seconds=5))
+        assert exc.value.is_transport_error is False
+        assert polls == 2
+
+
 def test_remote_job_wait_raises_on_failure():
     from lancedb.exceptions import JobFailedError
     from lancedb.index import BTree

python/src/error.rs

@@ -71,6 +71,7 @@ impl<T> PythonErrorExt<T> for std::result::Result<T, LanceError> {
                         request_id,
                         status_code.map(|s| s.as_u16()),
                         err.http_error_details(),
+                        is_transport_error(source.as_ref()),
                     )?;
 
                     if let Some(cause) = source.source() {
@@ -170,7 +171,14 @@ fn http_from_rust_error(
     let details = err
         .downcast_ref::<LanceError>()
         .and_then(LanceError::http_error_details);
-    let py_err = new_py_http_error(py, message, request_id, status_code, details)?;
+    let py_err = new_py_http_error(
+        py,
+        message,
+        request_id,
+        status_code,
+        details,
+        is_transport_error(err),
+    )?;
 
     // Reset the traceback since it doesn't provide additional information.
     let py_err = py_err.call_method1(intern!(py, "with_traceback"), (PyNone::get(py),))?;
@@ -189,6 +197,7 @@ fn new_py_http_error<'py>(
     request_id: &str,
     status_code: Option<u16>,
     details: Option<&HttpErrorDetails>,
+    is_transport_error: bool,
 ) -> PyResult<Bound<'py, PyAny>> {
     let error_code = details.map(|details| details.error_code.as_str());
     let resource_kind = details.map(|details| details.resource_kind.as_str());
@@ -205,5 +214,20 @@ fn new_py_http_error<'py>(
             resource_kind,
             resource_name,
             namespace_path,
+            is_transport_error,
         ))
 }
+
+// Preserve the typed transport classification across the Python boundary.
+// A missing status code alone also covers malformed job responses, which
+// waiting longer cannot repair.
+fn is_transport_error(err: &(dyn std::error::Error + 'static)) -> bool {
+    if let Some(error) = err.downcast_ref::<reqwest::Error>() {
+        return error.is_connect()
+            || error.is_timeout()
+            || error.is_request()
+            || error.is_body()
+            || error.is_decode();
+    }
+    err.source().is_some_and(is_transport_error)
+}