Handle backoff with new error

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:

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:

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.

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:

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)

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