From 3e0d1fe56be302567230436ca774a7527040c841 Mon Sep 17 00:00:00 2001 From: Kyle Consalus Date: Fri, 2 Oct 2026 18:04:44 -0700 Subject: [PATCH] fix(client): align task retry state with worker policy --- clients/python/src/taskbroker_client/retry.py | 6 +- clients/python/src/taskbroker_client/state.py | 7 +- .../taskbroker_client/worker/workerchild.py | 11 ++- clients/python/tests/worker/test_worker.py | 77 ++++++++++++++++++- 4 files changed, 92 insertions(+), 9 deletions(-) diff --git a/clients/python/src/taskbroker_client/retry.py b/clients/python/src/taskbroker_client/retry.py index 213248d5..49289292 100644 --- a/clients/python/src/taskbroker_client/retry.py +++ b/clients/python/src/taskbroker_client/retry.py @@ -59,7 +59,11 @@ def retry_task(exc: Exception | None = None, raise_on_no_retries: bool = True) - class Retry: - """Used with tasks to define the retry policy for a task""" + """Used with tasks to define the retry policy for a task. + + ``times`` is the total number of attempts, including the initial execution. + For example, ``Retry(times=3)`` allows the initial attempt plus two retries. + """ def __init__( self, diff --git a/clients/python/src/taskbroker_client/state.py b/clients/python/src/taskbroker_client/state.py index 9217262c..c96e0491 100644 --- a/clients/python/src/taskbroker_client/state.py +++ b/clients/python/src/taskbroker_client/state.py @@ -23,16 +23,13 @@ def current_task() -> CurrentTaskState | None: return _current_state.state -def set_current_task(activation: TaskActivation) -> None: - retry_state = activation.retry_state +def set_current_task(activation: TaskActivation, *, retries_remaining: bool) -> None: state = CurrentTaskState( id=activation.id, namespace=activation.namespace, taskname=activation.taskname, attempt=activation.retry_state.attempts, - # We subtract one, as attempts starts at 0, but `max_attempts` - # starts at 1. - retries_remaining=(retry_state.attempts < (retry_state.max_attempts - 1)), + retries_remaining=retries_remaining, processing_deadline_duration=activation.processing_deadline_duration, ) _current_state.state = state diff --git a/clients/python/src/taskbroker_client/worker/workerchild.py b/clients/python/src/taskbroker_client/worker/workerchild.py index 1f65ce0d..f62d8d0f 100644 --- a/clients/python/src/taskbroker_client/worker/workerchild.py +++ b/clients/python/src/taskbroker_client/worker/workerchild.py @@ -499,7 +499,14 @@ def check_task_future_completion( ) continue - set_current_task(inflight.activation) + retry = task_func.retry + set_current_task( + inflight.activation, + retries_remaining=( + retry is not None + and not retry.max_attempts_reached(inflight.activation.retry_state) + ), + ) next_state = TASK_ACTIVATION_STATUS_FAILURE # Use time.time() so we can measure against activation.received_at @@ -858,7 +865,7 @@ def _place_processing_result( # Don't send it on every task as this codepath is relatively # unoptimized on the broker side. max_attempts=( - task_func.retry._times + 1 + task_func.retry._times if task_func.retry and next_state == TASK_ACTIVATION_STATUS_RETRY else None ), diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 4a1febe6..0c042494 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -44,7 +44,7 @@ from taskbroker_client.canary import CANARY_TASK_NAME from taskbroker_client.constants import INTERNAL_NAMESPACE, CompressionType -from taskbroker_client.retry import NoRetriesRemainingError +from taskbroker_client.retry import NoRetriesRemainingError, Retry, RetryTaskError, retry_task from taskbroker_client.state import current_task from taskbroker_client.types import InflightTaskActivation, ProcessingResult from taskbroker_client.worker.childtiming import ( @@ -2393,6 +2393,81 @@ def test_child_process_retry_task() -> None: result = processed.get() assert result.task_id == RETRY_TASK.activation.id assert result.status == TASK_ACTIVATION_STATUS_RETRY + assert result.max_attempts == 2 + + +@pytest.mark.parametrize( + "times,retry_state,retries_remaining", + [ + pytest.param(3, RetryState(attempts=0, max_attempts=3), True, id="initial-attempt"), + pytest.param(3, RetryState(attempts=1, max_attempts=3), True, id="penultimate-attempt"), + pytest.param(3, RetryState(attempts=2, max_attempts=3), False, id="final-attempt"), + pytest.param(3, RetryState(attempts=2, max_attempts=4), False, id="inflated-limit"), + pytest.param(3, RetryState(attempts=1, max_attempts=2), True, id="stale-lower-limit"), + pytest.param(3, None, True, id="missing-retry-state"), + pytest.param(1, RetryState(attempts=0, max_attempts=3), False, id="single-attempt"), + pytest.param(None, RetryState(attempts=0, max_attempts=3), False, id="no-policy"), + ], +) +def test_child_process_retry_helper_uses_worker_policy( + times: int | None, + retry_state: RetryState | None, + retries_remaining: bool, + restore_signal_handlers: None, +) -> None: + from examples.tasks import retry_task as example_task + + observed_retries_remaining: list[bool] = [] + observed_errors: list[type[RetryTaskError]] = [] + + def request_retry() -> None: + current = current_task() + assert current is not None + observed_retries_remaining.append(current.retries_remaining) + try: + retry_task() + except RetryTaskError as err: + observed_errors.append(type(err)) + raise + + activation = TaskActivation() + activation.CopyFrom(RETRY_TASK.activation) + if retry_state is not None: + activation.retry_state.CopyFrom(retry_state) + + todo: queue.Queue[InflightTaskActivation] = queue.Queue() + processed: queue.Queue[ProcessingResult] = queue.Queue() + todo.put( + InflightTaskActivation(host="localhost:50051", receive_timestamp=0, activation=activation) + ) + + with ( + mock.patch.object( + example_task, "_retry", Retry(times=times) if times is not None else None + ), + mock.patch.object(example_task, "_func", request_retry), + ): + child_process( + "examples.app:app", + todo, + processed, + Event(), + max_task_count=1, + processing_pool_name="test", + process_type="fork", + skip_awaiting_futures=False, + future_checking_frequency=0.1, + ) + + assert observed_retries_remaining == [retries_remaining] + assert observed_errors == [RetryTaskError if retries_remaining else NoRetriesRemainingError] + result = processed.get(timeout=1) + assert result.task_id == activation.id + assert result.status == ( + TASK_ACTIVATION_STATUS_RETRY if retries_remaining else TASK_ACTIVATION_STATUS_FAILURE + ) + assert result.max_attempts == (times if retries_remaining else None) + assert current_task() is None @mock.patch("taskbroker_client.worker.workerchild.logger")