diff --git a/src/flyte/_internal/controllers/remote/_action.py b/src/flyte/_internal/controllers/remote/_action.py index ffd2234e2..731bbbe03 100644 --- a/src/flyte/_internal/controllers/remote/_action.py +++ b/src/flyte/_internal/controllers/remote/_action.py @@ -51,6 +51,7 @@ class Action: client_err: Exception | None = None # This error is set when something goes wrong in the controller. cache_key: str | None = None # None means no caching, otherwise it is the version of the cache. condition_output: literals_pb2.Literal | None = None # Output Literal for condition actions (set from ActionUpdate) + resource_exhausted_retries: int = 0 @property def name(self) -> str: diff --git a/src/flyte/_internal/controllers/remote/_core.py b/src/flyte/_internal/controllers/remote/_core.py index 9a6c6922c..372778322 100644 --- a/src/flyte/_internal/controllers/remote/_core.py +++ b/src/flyte/_internal/controllers/remote/_core.py @@ -501,6 +501,8 @@ async def _bg_launch(self, action: Action): if e.code == Code.ALREADY_EXISTS: logger.info(f"Action {action.name} already exists, continuing to monitor.") return + if e.code == Code.RESOURCE_EXHAUSTED: + raise flyte.errors.ResourceExhaustedError(e.message) from e if e.code == Code.ABORTED: # The run was aborted; engine will auto-abort other in-flight actions. # Surface as a system error — outer handler in _bg_run wraps and exits. @@ -555,6 +557,22 @@ async def _bg_run(self, worker_id: str): try: try: await self._bg_process(action) + except flyte.errors.ResourceExhaustedError as e: + action.resource_exhausted_retries += 1 + backoff = min( + self._min_backoff_on_err * (2 ** min(action.resource_exhausted_retries - 1, 20)), + self._max_backoff_on_err, + ) + logger.warning( + "Resource exhausted for action %s; retrying in %.1fs (attempt %d): %s", + action.name, + backoff, + action.resource_exhausted_retries, + e, + ) + await asyncio.sleep(backoff) + if self._running and not action.is_terminal(): + await self._shared_queue.put(action) except flyte.errors.SlowDownError as e: action.retries += 1 if action.retries > self._max_retries: diff --git a/src/flyte/errors.py b/src/flyte/errors.py index 4ed1f7332..6b34c500f 100644 --- a/src/flyte/errors.py +++ b/src/flyte/errors.py @@ -277,6 +277,10 @@ def __init__(self, message: str): super().__init__("SlowDownError", message, "user") +class ResourceExhaustedError(SlowDownError): + pass + + class OnlyAsyncIOSupportedError(RuntimeUserError): """ This error is raised when the user tries to use sync IO in an async task. diff --git a/src/flyte/remote/_client/auth/_interceptors/retry.py b/src/flyte/remote/_client/auth/_interceptors/retry.py index 62d6983c2..3611a5b87 100644 --- a/src/flyte/remote/_client/auth/_interceptors/retry.py +++ b/src/flyte/remote/_client/auth/_interceptors/retry.py @@ -8,27 +8,20 @@ from flyte._logging import logger -RETRYABLE_CODES = frozenset({Code.UNAVAILABLE, Code.RESOURCE_EXHAUSTED, Code.INTERNAL}) +RETRYABLE_CODES = frozenset({Code.UNAVAILABLE, Code.INTERNAL}) def _log_retry(ctx, e: ConnectError, delay: float, attempt: int, max_attempts: int) -> None: method = getattr(getattr(ctx, "method", None), "name", "rpc") - if e.code == Code.RESOURCE_EXHAUSTED: - # The server explicitly asked us to back off (e.g. a queue at max - # depth); surface why the call is pausing instead of sleeping silently. - logger.warning( - "%s rejected: %s; retrying in %.1fs (attempt %d/%d)", method, e.message, delay, attempt + 1, max_attempts - ) - else: - logger.debug( - "%s failed (code=%s): %s; retrying in %.1fs (attempt %d/%d)", - method, - e.code, - e.message, - delay, - attempt + 1, - max_attempts, - ) + logger.debug( + "%s failed (code=%s): %s; retrying in %.1fs (attempt %d/%d)", + method, + e.code, + e.message, + delay, + attempt + 1, + max_attempts, + ) class RetryUnaryInterceptor: diff --git a/tests/flyte/remote/test_connectrpc_interceptors.py b/tests/flyte/remote/test_connectrpc_interceptors.py index 48b66ec76..85d36abbc 100644 --- a/tests/flyte/remote/test_connectrpc_interceptors.py +++ b/tests/flyte/remote/test_connectrpc_interceptors.py @@ -569,18 +569,15 @@ async def test_retries_on_unavailable(self): assert call_next.call_count == 2 @pytest.mark.asyncio - async def test_retries_on_resource_exhausted(self): + async def test_does_not_retry_resource_exhausted(self): interceptor = RetryUnaryInterceptor(max_attempts=3, initial_backoff=0.001) - call_next = AsyncMock( - side_effect=[ - ConnectError(Code.RESOURCE_EXHAUSTED, "exhausted"), - "ok", - ] - ) + call_next = AsyncMock(side_effect=ConnectError(Code.RESOURCE_EXHAUSTED, "exhausted")) ctx, _ = _make_ctx_mock() - result = await interceptor.intercept_unary(call_next, "req", ctx) - assert result == "ok" + with pytest.raises(ConnectError) as exc_info: + await interceptor.intercept_unary(call_next, "req", ctx) + assert exc_info.value.code == Code.RESOURCE_EXHAUSTED + assert call_next.call_count == 1 @pytest.mark.asyncio async def test_retries_on_internal(self): @@ -649,28 +646,6 @@ async def mock_sleep(duration): assert 1.0 <= sleep_durations[1] < 3.0 # base=2.0 assert 2.0 <= sleep_durations[2] < 6.0 # base=4.0 - @pytest.mark.asyncio - async def test_resource_exhausted_retry_logs_warning(self, caplog, monkeypatch): - # The "flyte" logger disables propagation; re-enable it so caplog sees records. - monkeypatch.setattr(logging.getLogger("flyte"), "propagate", True) - interceptor = RetryUnaryInterceptor(max_attempts=3, initial_backoff=0.001) - call_next = AsyncMock( - side_effect=[ - ConnectError(Code.RESOURCE_EXHAUSTED, 'queue "q" is at max depth (5), retry later'), - "ok", - ] - ) - ctx, _ = _make_ctx_mock() - - with caplog.at_level("WARNING", logger="flyte"): - result = await interceptor.intercept_unary(call_next, "req", ctx) - - assert result == "ok" - warnings = [r for r in caplog.records if r.levelname == "WARNING"] - assert len(warnings) == 1 - assert 'queue "q" is at max depth (5)' in warnings[0].getMessage() - assert "attempt 1/3" in warnings[0].getMessage() - @pytest.mark.asyncio async def test_unavailable_retry_logs_debug_only(self, caplog, monkeypatch): monkeypatch.setattr(logging.getLogger("flyte"), "propagate", True) diff --git a/tests/internal/controllers/test_remote_controller.py b/tests/internal/controllers/test_remote_controller.py index fd48c51f6..a3a8b669e 100644 --- a/tests/internal/controllers/test_remote_controller.py +++ b/tests/internal/controllers/test_remote_controller.py @@ -610,6 +610,55 @@ async def fake_bg_process(action): informer.fire_completion_event.assert_awaited_once_with("A") +@pytest.mark.asyncio +async def test_bg_run_retries_resource_exhausted_without_consuming_retry_budget(): + from aiolimiter import AsyncLimiter + from connectrpc.code import Code + from connectrpc.errors import ConnectError + + controller = object.__new__(Controller) + controller._running = True + controller._shared_queue = asyncio.Queue() + controller._max_retries = 1 + controller._min_backoff_on_err = 0.1 + controller._max_backoff_on_err = 0.1 + controller._rate_limiter = AsyncLimiter(100, 1.0) + controller._enqueue_timeout = 0.01 + controller._consecutive_launch_timeouts = 0 + controller._actions_service = AsyncMock() + + action = Action( + parent_action_name="parent", + action_id=identifier_pb2.ActionIdentifier( + name="child", + run=identifier_pb2.RunIdentifier(name="run"), + ), + type="trace", + retries=controller._max_retries, + ) + await controller._shared_queue.put(action) + + enqueue_attempts = 0 + + async def enqueue(*_args, **_kwargs): + nonlocal enqueue_attempts + enqueue_attempts += 1 + if enqueue_attempts <= 3: + raise ConnectError(Code.RESOURCE_EXHAUSTED, "exhausted") + controller._running = False + + controller._actions_service.enqueue.side_effect = enqueue + + with patch("flyte._internal.controllers.remote._core.asyncio.sleep", new_callable=AsyncMock) as sleep: + await controller._bg_run(worker_id="w1") + + assert enqueue_attempts == 4 + assert action.resource_exhausted_retries == 3 + assert action.retries == controller._max_retries + assert action.client_err is None + assert sleep.await_count == 3 + + @pytest.mark.asyncio async def test_record_trace_uses_task_action_when_in_trace_scope(): """When tctx.action has been swapped by @trace, record_trace must submit with