From ad4eeba49d7e317904d0fad06e2c0ffa5a1ff315 Mon Sep 17 00:00:00 2001 From: Kayvan Zahiri <123409429+Kayvan-Zahiri@users.noreply.github.com> Date: Wed, 21 Oct 2026 21:07:48 -0700 Subject: [PATCH] Fix task result retries after transient errors --- src/parallel/lib/_parsing/_task_run_result.py | 53 ++++++++++-- src/parallel/lib/_time.py | 55 +----------- tests/test_task_result_retry.py | 86 +++++++++++++++++++ 3 files changed, 134 insertions(+), 60 deletions(-) create mode 100644 tests/test_task_result_retry.py diff --git a/src/parallel/lib/_parsing/_task_run_result.py b/src/parallel/lib/_parsing/_task_run_result.py index bd3c02f..dc6c3f8 100644 --- a/src/parallel/lib/_parsing/_task_run_result.py +++ b/src/parallel/lib/_parsing/_task_run_result.py @@ -1,8 +1,8 @@ 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 @@ -10,6 +10,7 @@ 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) diff --git a/src/parallel/lib/_time.py b/src/parallel/lib/_time.py index 80a2306..881125c 100644 --- a/src/parallel/lib/_time.py +++ b/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) diff --git a/tests/test_task_result_retry.py b/tests/test_task_result_retry.py new file mode 100644 index 0000000..96d6b80 --- /dev/null +++ b/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