diff --git a/modelq/app/base.py b/modelq/app/base.py index b786361..bec6d07 100644 --- a/modelq/app/base.py +++ b/modelq/app/base.py @@ -153,6 +153,7 @@ def __init__( webhook_url: Optional[str] = None, # Optional webhook for error logging requeue_threshold : Optional[int] = None , delay_seconds: int = 30, + retry_task_max_attempts: int = 3, # Cap on RetryTaskException re-queues task_history_retention: Optional[int] = None, # Configurable history retention (default 24h) task_ttl: Optional[int] = None, # Configurable task TTL (default 24h) inactive_if_worker_boot_fail: bool = False, # Mark worker as unhealthy if before_worker_boot fails @@ -194,6 +195,7 @@ def __init__( self.webhook_url = webhook_url self.requeue_threshold = requeue_threshold self.delay_seconds = delay_seconds + self.retry_task_max_attempts = retry_task_max_attempts self.task_history_retention = task_history_retention or self.TASK_HISTORY_RETENTION self.task_ttl = task_ttl or self.TASK_TTL self.inactive_if_worker_boot_fail = inactive_if_worker_boot_fail @@ -1035,10 +1037,41 @@ def process_task(self, task: Task) -> None: logger.info(f"Task {task.task_name} completed successfully.") except RetryTaskException as e: - logger.warning(f"Task {task.task_name} requested retry: {e}") - new_task_dict = task.to_dict() - new_task_dict["payload"] = task.original_payload - self.enqueue_delayed_task(new_task_dict, delay_seconds=self.delay_seconds) + # A RetryTaskException used to re-queue unconditionally, with no + # attempt counter. Any task whose failure is permanent -- an image + # URL that has expired or been deleted, for example -- therefore + # looped forever, re-entering the queue every delay_seconds and + # burning worker capacity that never produced a result. + # Bound it. The counter travels on the payload so it survives the + # round trip through Redis, and uses a reserved key so it cannot + # collide with the caller's own "retries" budget used elsewhere. + payload = task.original_payload if isinstance(task.original_payload, dict) else {} + attempts = int(payload.get("_retry_attempts", 0) or 0) + 1 + max_attempts = max(0, int(getattr(self, "retry_task_max_attempts", 3))) + + if attempts > max_attempts: + task.status = "failed" + task.result = ( + f"Retry limit reached after {max_attempts} attempt(s): {e}" + ) + self._store_final_task_state(task, success=False, error=e) + self.log_task_error_to_file(task, e) + self.check_middleware("on_error", task=task, error=e) + self.post_error_to_webhook(task, e) + logger.error( + f"Task {task.task_name} exhausted its retry budget " + f"({max_attempts}); marking failed: {e}" + ) + else: + logger.warning( + f"Task {task.task_name} requested retry " + f"({attempts}/{max_attempts}): {e}" + ) + new_task_dict = task.to_dict() + new_task_dict["payload"] = task.original_payload + if isinstance(new_task_dict.get("payload"), dict): + new_task_dict["payload"]["_retry_attempts"] = attempts + self.enqueue_delayed_task(new_task_dict, delay_seconds=self.delay_seconds) except Exception as e: # Mark as failed task.status = "failed" diff --git a/tests/test_retry_budget.py b/tests/test_retry_budget.py new file mode 100644 index 0000000..dedbe15 --- /dev/null +++ b/tests/test_retry_budget.py @@ -0,0 +1,89 @@ +"""Regression tests for the RetryTaskException retry budget. + +A RetryTaskException used to be re-queued unconditionally, with no attempt +counter. Any task whose failure was permanent -- an expired or deleted source +URL, for example -- looped forever, re-entering the queue every delay_seconds +and consuming worker capacity that could never produce a result. These tests +pin the bound. +""" + +import fakeredis +import pytest + +from modelq import ModelQ +from modelq.app.tasks.base import Task +from modelq.exceptions import RetryTaskException + + +@pytest.fixture +def mock_redis(): + return fakeredis.FakeStrictRedis() + + +def _register_doomed(mq): + """Register a task that always asks to be retried.""" + def doomed(*args, **kwargs): + raise RetryTaskException("source image is gone for good") + + # process_task resolves the callable via getattr(self, task.task_name) + mq.allowed_tasks.add("doomed") + setattr(mq, "doomed", doomed) + + +def _queued_delayed(mq): + """Number of tasks sitting in the delayed set.""" + return mq.redis_client.zcard("delayed_tasks") + + +def test_permanently_failing_task_stops_retrying(mock_redis): + mq = ModelQ(redis_client=mock_redis, retry_task_max_attempts=2) + _register_doomed(mq) + + task = Task(task_name="doomed", payload={"data": {"args": [], "kwargs": {}}, "init_image": "https://gone.example/x.png"}) + + # Attempt 1 -> re-queued, counter starts. + mq.process_task(task) + assert task.original_payload["_retry_attempts"] == 1 + assert task.status != "failed" + + # Attempt 2 -> re-queued, counter advances. + mq.process_task(task) + assert task.original_payload["_retry_attempts"] == 2 + assert task.status != "failed" + + # Attempt 3 exceeds the budget -> the task is failed, not re-queued. + before = _queued_delayed(mq) + mq.process_task(task) + assert task.status == "failed" + assert "Retry limit reached" in str(task.result) + assert _queued_delayed(mq) == before, "must not re-queue once the budget is spent" + + +def test_reserved_counter_does_not_consume_the_callers_retries_budget(mock_redis): + mq = ModelQ(redis_client=mock_redis, retry_task_max_attempts=5) + _register_doomed(mq) + + task = Task(task_name="doomed", payload={"data": {"args": [], "kwargs": {}}, "retries": 4}) + mq.process_task(task) + + # "retries" belongs to the TaskProcessingError path and must be untouched. + assert task.original_payload["retries"] == 4 + assert task.original_payload["_retry_attempts"] == 1 + + +def test_zero_budget_fails_on_first_retry_request(mock_redis): + mq = ModelQ(redis_client=mock_redis, retry_task_max_attempts=0) + _register_doomed(mq) + + task = Task(task_name="doomed", payload={"data": {"args": [], "kwargs": {}}}) + mq.process_task(task) + + assert task.status == "failed" + assert "Retry limit reached" in str(task.result) + + +def test_default_budget_is_finite(mock_redis): + """The old behaviour was unbounded; the default must not be.""" + mq = ModelQ(redis_client=mock_redis) + assert isinstance(mq.retry_task_max_attempts, int) + assert 0 < mq.retry_task_max_attempts < 100