Fix task result retries after transient errors

src/parallel/lib/_parsing/_task_run_result.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
-from typing import Any, Type, Union, Callable, Coroutine, cast
+import time
+from typing import Any, Type, Union, Callable, NoReturn, Coroutine, cast
 
-from .._time import timeout_retry_context
 from ...types import TaskRunResult, ParsedTaskRunResult
 from ..._types import Omit
 from ..._utils import is_str, is_given
 from ..._compat import model_parse, model_parse_json
 from ..._models import construct_type_unchecked
 from .._pydantic import is_basemodel_type
 from ._task_spec import is_output_schema_param
+from ..._exceptions import APIStatusError, APITimeoutError
 from ...types.task_spec_param import OutputT, OutputSchema
 
 
@@ -20,8 +21,17 @@ def wait_for_result(
     callable: Callable[[str, float], TaskRunResult | ParsedTaskRunResult[OutputT]],
 ) -> TaskRunResult | ParsedTaskRunResult[OutputT]:
     """Wait for a task run to complete within the given timeout."""
-    with timeout_retry_context(run_id, deadline):
-        return callable(run_id, deadline)
+    exc: Exception | None = None
+    while time.monotonic() < deadline:
+        try:
+            return callable(run_id, deadline)
+        except APITimeoutError as error:
+            exc = error
+        except APIStatusError as error:
+            if not _is_retryable_error(error.status_code):
+                raise
+            exc = error
+    _raise_timeout(run_id, exc)
 
 
 async def wait_for_result_async(
@@ -31,8 +41,17 @@ async def wait_for_result_async(
     callable: Callable[[str, float], Coroutine[Any, Any, TaskRunResult | ParsedTaskRunResult[OutputT]]],
 ) -> TaskRunResult | ParsedTaskRunResult[OutputT]:
     """Wait for a task run to complete within the given timeout."""
-    with timeout_retry_context(run_id, deadline):
-        return await callable(run_id, deadline)
+    exc: Exception | None = None
+    while time.monotonic() < deadline:
+        try:
+            return await callable(run_id, deadline)
+        except APITimeoutError as error:
+            exc = error
+        except APIStatusError as error:
+            if not _is_retryable_error(error.status_code):
+                raise
+            exc = error
+    _raise_timeout(run_id, exc)
 
 
 def task_run_result_parser(
@@ -97,3 +116,25 @@ def _parse_content(response_format: type[OutputT], content: str | object) -> Out
             return cast(OutputT, model_parse(response_format, content))
 
     raise TypeError(f"Unable to automatically parse response format type {response_format}")
+
+
+def _raise_timeout(run_id: str, exc: Union[Exception, None]) -> NoReturn:
+    raise TimeoutError(f"Fetching task run result for run id {run_id} timed out.") from exc
+
+
+def _is_retryable_error(status_code: int) -> bool:
+    """Determine if an error is retryable.
+
+    We retry the following HTTP status codes within the SDK:
+    - 408 (Request Timeout): The server timed out waiting for the request
+    - 503 (Service Unavailable): The server is temporarily unable to handle the request
+    - 504 (Gateway Timeout): The gateway server timed out
+
+    These errors typically indicate temporary issues with the server or network
+    that may resolve upon retry. We don't include 429 (Too Many Requests) as this
+    indicates rate limiting, which requires backing off rather than immediate retries.
+
+    Note: This is a low-level retry mechanism within the SDK. Customers may want to
+    implement their own retry logic at the application level for other error types.
+    """
+    return status_code in (408, 503, 504)

src/parallel/lib/_time.py

@@ -1,12 +1,9 @@
-import time
-import contextlib
-from typing import Union, Iterator, NoReturn
+from typing import Union
 
 import httpx
 
 from .._types import NotGiven
 from .._constants import DEFAULT_EXECUTE_TIMEOUT_SECONDS
-from .._exceptions import APIStatusError, APITimeoutError
 from .._utils._utils import is_given
 
 
@@ -22,53 +19,3 @@ def prepare_timeout_float(timeout: Union[float, httpx.Timeout, None, NotGiven])
         return DEFAULT_EXECUTE_TIMEOUT_SECONDS
 
     return timeout
-
-
-def _raise_timeout(run_id: str, exc: Union[Exception, None]) -> NoReturn:
-    raise TimeoutError(f"Fetching task run result for run id {run_id} timed out.") from exc
-
-
-def _is_retryable_error(status_code: int) -> bool:
-    """Determine if an error is retryable.
-
-    We retry the following HTTP status codes within the SDK:
-    - 408 (Request Timeout): The server timed out waiting for the request
-    - 503 (Service Unavailable): The server is temporarily unable to handle the request
-    - 504 (Gateway Timeout): The gateway server timed out
-
-    These errors typically indicate temporary issues with the server or network
-    that may resolve upon retry. We don't include 429 (Too Many Requests) as this
-    indicates rate limiting, which requires backing off rather than immediate retries.
-
-    Note: This is a low-level retry mechanism within the SDK. Customers may want to
-    implement their own retry logic at the application level for other error types.
-    """
-    return status_code in (408, 503, 504)
-
-
-@contextlib.contextmanager
-def timeout_retry_context(run_id: str, deadline: float) -> Iterator[None]:
-    """Context manager for handling timeouts and retries when fetching task run results.
-
-    Args:
-        run_id: The ID of the task run
-        deadline: The absolute time (monotonic) by which the operation must complete
-
-    Raises:
-        TimeoutError: If the deadline is reached
-        APIStatusError: For non-timeout API errors
-    """
-    exc: Union[Exception, None] = None
-    while time.monotonic() < deadline:
-        try:
-            yield
-            return
-        except APITimeoutError as e:
-            exc = e
-            continue
-        except APIStatusError as e:
-            if _is_retryable_error(e.status_code):
-                exc = e
-                continue
-            raise
-    _raise_timeout(run_id, exc)

tests/test_task_result_retry.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import time
+from unittest.mock import Mock
+
+import httpx
+import pytest
+
+from parallel import Parallel, AsyncParallel, APIStatusError
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("failure", [408, 503, 504, "timeout"])
+async def test_result_retries_after_transient_failure(asynchronous: bool, failure: int | str) -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        if request.method == "POST":
+            return httpx.Response(200, json={"run_id": "run_test"})
+        calls += 1
+        if calls == 1:
+            if failure == "timeout":
+                raise httpx.ReadTimeout("temporary", request=request)
+            return httpx.Response(int(failure), json={"error": "temporary"})
+        return httpx.Response(200, json={"run_id": "run_test", "output": {"type": "text", "content": "done"}})
+
+    if asynchronous:
+        async with AsyncParallel(
+            api_key="test", max_retries=0, http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler))
+        ) as client:
+            result = await client.task_run.execute(input="test", processor="lite", timeout=10)
+    else:
+        with Parallel(
+            api_key="test", max_retries=0, http_client=httpx.Client(transport=httpx.MockTransport(handler))
+        ) as client:
+            result = client.task_run.execute(input="test", processor="lite", timeout=10)
+    assert result.output.content == "done"
+    assert calls == 2
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_non_retryable_error_propagates(asynchronous: bool) -> None:
+    handler = Mock(return_value=httpx.Response(429, json={"error": "rate limited"}))
+    with pytest.raises(APIStatusError) as exc:
+        if asynchronous:
+            async with AsyncParallel(
+                api_key="test", max_retries=0, http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler))
+            ) as client:
+                await client.task_run._wait_for_result(run_id="run_test", deadline=time.monotonic() + 10)
+        else:
+            with Parallel(
+                api_key="test", max_retries=0, http_client=httpx.Client(transport=httpx.MockTransport(handler))
+            ) as client:
+                client.task_run._wait_for_result(run_id="run_test", deadline=time.monotonic() + 10)
+    assert exc.value.status_code == 429
+    assert handler.call_count == 1
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("expired", [False, True])
+async def test_deadline_and_last_error(monkeypatch: pytest.MonkeyPatch, asynchronous: bool, expired: bool) -> None:
+    from types import SimpleNamespace
+
+    from parallel.lib._parsing import _task_run_result
+
+    clock = Mock(side_effect=[2] if expired else [0, 2])
+    monkeypatch.setattr(_task_run_result, "time", SimpleNamespace(monotonic=clock))
+    handler = Mock(return_value=httpx.Response(503, json={"error": "temporary"}))
+    with pytest.raises(TimeoutError, match="run_test") as exc:
+        if asynchronous:
+            async with AsyncParallel(
+                api_key="test", max_retries=0, http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler))
+            ) as client:
+                await client.task_run._wait_for_result(run_id="run_test", deadline=1)
+        else:
+            with Parallel(
+                api_key="test", max_retries=0, http_client=httpx.Client(transport=httpx.MockTransport(handler))
+            ) as client:
+                client.task_run._wait_for_result(run_id="run_test", deadline=1)
+    assert handler.call_count == (0 if expired else 1)
+    if expired:
+        assert exc.value.__cause__ is None
+    else:
+        assert isinstance(exc.value.__cause__, APIStatusError)
+        assert exc.value.__cause__.status_code == 503