From 62b8c559b66288d822da4ee0713927a756a17c0d Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 28 Aug 2026 17:54:57 -0700 Subject: [PATCH 1/3] handle backoff Signed-off-by: Yee Hing Tong --- .../_internal/controllers/remote/_action.py | 1 + .../_internal/controllers/remote/_core.py | 19 +++++++ .../_client/auth/_interceptors/retry.py | 2 +- .../remote/test_connectrpc_interceptors.py | 15 +++--- .../controllers/test_remote_controller.py | 49 +++++++++++++++++++ 5 files changed, 76 insertions(+), 10 deletions(-) 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..e6845e809 100644 --- a/src/flyte/_internal/controllers/remote/_core.py +++ b/src/flyte/_internal/controllers/remote/_core.py @@ -401,6 +401,23 @@ async def _bg_cancel_action(self, action: Action): if informer: await informer.fire_completion_event(action.name) + async def _bg_handle_resource_exhausted(self, action: Action, error: ConnectError) -> None: + 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, + error.message, + ) + await asyncio.sleep(backoff) + if self._running and not action.is_terminal(): + await self._shared_queue.put(action) + async def _bg_launch(self, action: Action): """ Attempt to launch an action. @@ -501,6 +518,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: + return await self._bg_handle_resource_exhausted(action, 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. diff --git a/src/flyte/remote/_client/auth/_interceptors/retry.py b/src/flyte/remote/_client/auth/_interceptors/retry.py index 0519aa1c9..c7c45f221 100644 --- a/src/flyte/remote/_client/auth/_interceptors/retry.py +++ b/src/flyte/remote/_client/auth/_interceptors/retry.py @@ -6,7 +6,7 @@ from connectrpc.code import Code from connectrpc.errors import ConnectError -RETRYABLE_CODES = frozenset({Code.UNAVAILABLE, Code.RESOURCE_EXHAUSTED, Code.INTERNAL}) +RETRYABLE_CODES = frozenset({Code.UNAVAILABLE, Code.INTERNAL}) class RetryUnaryInterceptor: diff --git a/tests/flyte/remote/test_connectrpc_interceptors.py b/tests/flyte/remote/test_connectrpc_interceptors.py index 5df92fa57..46fb93da7 100644 --- a/tests/flyte/remote/test_connectrpc_interceptors.py +++ b/tests/flyte/remote/test_connectrpc_interceptors.py @@ -568,18 +568,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): 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 From 1d8f699b457f4773a0c146135fa7368ef5e5663a Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Tue, 1 Sep 2026 08:27:47 -0700 Subject: [PATCH 2/3] resource exhausted Signed-off-by: Yee Hing Tong --- examples/cluster_drain/README.md | 26 +++++++ examples/cluster_drain/drain_workloads.py | 74 +++++++++++++++++++ examples/cluster_drain/resource_exhausted.py | 50 +++++++++++++ .../_internal/controllers/remote/_core.py | 35 +++++---- src/flyte/errors.py | 4 + 5 files changed, 171 insertions(+), 18 deletions(-) create mode 100644 examples/cluster_drain/README.md create mode 100644 examples/cluster_drain/drain_workloads.py create mode 100644 examples/cluster_drain/resource_exhausted.py diff --git a/examples/cluster_drain/README.md b/examples/cluster_drain/README.md new file mode 100644 index 000000000..9ad79a4ca --- /dev/null +++ b/examples/cluster_drain/README.md @@ -0,0 +1,26 @@ +# cluster_drain + +Workloads for the leasor cluster-draining staging regimen +(runbook: https://claude.ai/code/artifact/c9a04c77-50f8-411a-b5c9-8a17a8a81019). + +All tasks live in `drain_workloads.py` and take the queue from the command line. `A` below is the cluster being +drained; `Q` is a queue that routes to both `A` and `B` (wildcard or both pinned); `A-queue` is A's co-named queue. + +| Step | Command | +|---|---| +| T0 / T1 placement probe, after every `--activate` | `flyte run --queue Q drain_workloads.py quick` | +| T1–T5 long runner | `flyte run --queue Q drain_workloads.py sleep_for --seconds 1800` | +| T6 leasor restart mid-drain | `flyte run --queue Q drain_workloads.py fan_out --n 50 --seconds 300` | +| T7 co-named queue, terminal fail | `flyte run --queue A-queue drain_workloads.py fan_out --n 3 --seconds 1800` | +| T7 parked action | `flyte run --queue A-queue drain_workloads.py sleep_for --seconds 60` after the drain call | + +What to read in the task logs: + +- `sleep_for` prints `alive Ns` with the pod hostname every 10s. After a force drain on a multi-cluster queue the + same action reappears with a new hostname (attempt 2) — the old pod's lines stop. +- On the co-named queue nothing reappears; the run ends FAILED with `queue A-queue routes to no other cluster`. +- `fan_out` prints the per-host child count at the end; on T6 it should sum to `--n` with every child on B. + +To land work on `A` specifically for T1–T5, use a queue that routes only to `A` (its co-named queue) for the +*first* placement but expect T2/T3 to then fail it terminally — or, simpler, pause B's leaseworker while submitting +so `Q` places on `A`, then resume it before the drain so retries have somewhere to go. diff --git a/examples/cluster_drain/drain_workloads.py b/examples/cluster_drain/drain_workloads.py new file mode 100644 index 000000000..da590b001 --- /dev/null +++ b/examples/cluster_drain/drain_workloads.py @@ -0,0 +1,74 @@ +""" +Workloads for the cluster-draining staging regimen (leasor PRs #17847 / #17920 / #17940). + +Every task prints the pod hostname and a heartbeat line so leasor-side resets +are visible from the task logs alone: after a force drain the same action shows +up again with a new hostname on another cluster (system retry), while a task +whose queue routes only to the drained cluster never reappears (terminal FAIL). + +Pick the queue on the command line; the tasks don't pin one: + + flyte run --queue examples/cluster_drain/drain_workloads.py [--arg value] +""" + +import asyncio +import socket +import time + +import flyte + +env = flyte.TaskEnvironment( + name="cluster-drain", + resources=flyte.Resources(cpu=1, memory="256Mi"), +) + + +def _where() -> str: + return socket.gethostname() + + +@env.task +async def quick() -> str: + """Placement probe: lands, sleeps 5s, reports where it ran (T0, T1, and after every --activate).""" + host = _where() + print(f"quick: running on {host}") + await asyncio.sleep(5) + return host + + +@env.task +async def sleep_for(seconds: int = 1800) -> str: + """Long runner for T1-T5 and T7. Heartbeats every 10s so a reset shows as a hostname change in the logs.""" + host = _where() + start = time.monotonic() + print(f"sleep_for: started on {host}, will run {seconds}s") + while (elapsed := time.monotonic() - start) < seconds: + await asyncio.sleep(min(10, seconds - elapsed)) + print(f"sleep_for: {host} alive {int(time.monotonic() - start)}s/{seconds}s") + print(f"sleep_for: finished on {host}") + return host + + +@env.task +async def fan_out(n: int = 50, seconds: int = 300) -> int: + """ + Parent with n concurrent sleep_for children in the same queue. + + T6: --n 50 --seconds 300, then force-drain and restart the owning leasor shard; + every child should end on attempt 2, never 3. + T7: --n 3 --seconds 1800 on the co-named queue; the parent and all children must + FAIL (not requeue) and the run must end, with no child left Unassigned. + """ + print(f"fan_out: parent on {_where()}, spawning {n} children x {seconds}s") + hosts = await asyncio.gather(*(sleep_for(seconds) for _ in range(n))) + by_host: dict[str, int] = {} + for h in hosts: + by_host[h] = by_host.get(h, 0) + 1 + print(f"fan_out: children finished on {by_host}") + return len(hosts) + + +if __name__ == "__main__": + flyte.init_from_config() + run = flyte.run(quick) + print(run.name, run.url) diff --git a/examples/cluster_drain/resource_exhausted.py b/examples/cluster_drain/resource_exhausted.py new file mode 100644 index 000000000..fb4e817f5 --- /dev/null +++ b/examples/cluster_drain/resource_exhausted.py @@ -0,0 +1,50 @@ +"""Exercise RESOURCE_EXHAUSTED handling against a depth-limited queue. + +Run with a queue configured with run_concurrency=1, action_concurrency=2, +and depth=5: + + flyte run --queue examples/cluster_drain/resource_exhausted.py main +""" + +import asyncio +import socket +from datetime import datetime, timezone + +import flyte + +env = flyte.TaskEnvironment( + name="resource-exhausted", + resources=flyte.Resources(cpu=1, memory="256Mi"), +) + + +@env.task +async def hold(i: int, seconds: int) -> int: + host = socket.gethostname() + started = datetime.now(timezone.utc).isoformat(timespec="seconds") + print(f"hold {i}: started on {host} at {started}", flush=True) + await asyncio.sleep(seconds) + finished = datetime.now(timezone.utc).isoformat(timespec="seconds") + print(f"hold {i}: finished on {host} at {finished}", flush=True) + return i + + +@env.task +async def main(count: int = 8, seconds: int = 5) -> list[int]: + tasks: list[asyncio.Task[int]] = [] + for i in range(count): + tasks.append(asyncio.create_task(hold(i, seconds))) + await asyncio.sleep(0.5) + results = await asyncio.gather(*tasks) + expected = list(range(count)) + if results != expected: + raise AssertionError(f"expected {expected}, got {results}") + print(f"all {count} actions completed", flush=True) + return results + + +if __name__ == "__main__": + flyte.init_from_config() + run = flyte.run(main, count=8, seconds=5) + print(run.name, run.url) + run.wait() diff --git a/src/flyte/_internal/controllers/remote/_core.py b/src/flyte/_internal/controllers/remote/_core.py index e6845e809..372778322 100644 --- a/src/flyte/_internal/controllers/remote/_core.py +++ b/src/flyte/_internal/controllers/remote/_core.py @@ -401,23 +401,6 @@ async def _bg_cancel_action(self, action: Action): if informer: await informer.fire_completion_event(action.name) - async def _bg_handle_resource_exhausted(self, action: Action, error: ConnectError) -> None: - 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, - error.message, - ) - await asyncio.sleep(backoff) - if self._running and not action.is_terminal(): - await self._shared_queue.put(action) - async def _bg_launch(self, action: Action): """ Attempt to launch an action. @@ -519,7 +502,7 @@ async def _bg_launch(self, action: Action): logger.info(f"Action {action.name} already exists, continuing to monitor.") return if e.code == Code.RESOURCE_EXHAUSTED: - return await self._bg_handle_resource_exhausted(action, e) + 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. @@ -574,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. From beef8007d0a6c6352e6d44988bfa499f5a11b9a6 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Tue, 1 Sep 2026 08:28:24 -0700 Subject: [PATCH 3/3] remove bad examples Signed-off-by: Yee Hing Tong --- examples/cluster_drain/README.md | 26 ------- examples/cluster_drain/drain_workloads.py | 74 -------------------- examples/cluster_drain/resource_exhausted.py | 50 ------------- 3 files changed, 150 deletions(-) delete mode 100644 examples/cluster_drain/README.md delete mode 100644 examples/cluster_drain/drain_workloads.py delete mode 100644 examples/cluster_drain/resource_exhausted.py diff --git a/examples/cluster_drain/README.md b/examples/cluster_drain/README.md deleted file mode 100644 index 9ad79a4ca..000000000 --- a/examples/cluster_drain/README.md +++ /dev/null @@ -1,26 +0,0 @@ -# cluster_drain - -Workloads for the leasor cluster-draining staging regimen -(runbook: https://claude.ai/code/artifact/c9a04c77-50f8-411a-b5c9-8a17a8a81019). - -All tasks live in `drain_workloads.py` and take the queue from the command line. `A` below is the cluster being -drained; `Q` is a queue that routes to both `A` and `B` (wildcard or both pinned); `A-queue` is A's co-named queue. - -| Step | Command | -|---|---| -| T0 / T1 placement probe, after every `--activate` | `flyte run --queue Q drain_workloads.py quick` | -| T1–T5 long runner | `flyte run --queue Q drain_workloads.py sleep_for --seconds 1800` | -| T6 leasor restart mid-drain | `flyte run --queue Q drain_workloads.py fan_out --n 50 --seconds 300` | -| T7 co-named queue, terminal fail | `flyte run --queue A-queue drain_workloads.py fan_out --n 3 --seconds 1800` | -| T7 parked action | `flyte run --queue A-queue drain_workloads.py sleep_for --seconds 60` after the drain call | - -What to read in the task logs: - -- `sleep_for` prints `alive Ns` with the pod hostname every 10s. After a force drain on a multi-cluster queue the - same action reappears with a new hostname (attempt 2) — the old pod's lines stop. -- On the co-named queue nothing reappears; the run ends FAILED with `queue A-queue routes to no other cluster`. -- `fan_out` prints the per-host child count at the end; on T6 it should sum to `--n` with every child on B. - -To land work on `A` specifically for T1–T5, use a queue that routes only to `A` (its co-named queue) for the -*first* placement but expect T2/T3 to then fail it terminally — or, simpler, pause B's leaseworker while submitting -so `Q` places on `A`, then resume it before the drain so retries have somewhere to go. diff --git a/examples/cluster_drain/drain_workloads.py b/examples/cluster_drain/drain_workloads.py deleted file mode 100644 index da590b001..000000000 --- a/examples/cluster_drain/drain_workloads.py +++ /dev/null @@ -1,74 +0,0 @@ -""" -Workloads for the cluster-draining staging regimen (leasor PRs #17847 / #17920 / #17940). - -Every task prints the pod hostname and a heartbeat line so leasor-side resets -are visible from the task logs alone: after a force drain the same action shows -up again with a new hostname on another cluster (system retry), while a task -whose queue routes only to the drained cluster never reappears (terminal FAIL). - -Pick the queue on the command line; the tasks don't pin one: - - flyte run --queue examples/cluster_drain/drain_workloads.py [--arg value] -""" - -import asyncio -import socket -import time - -import flyte - -env = flyte.TaskEnvironment( - name="cluster-drain", - resources=flyte.Resources(cpu=1, memory="256Mi"), -) - - -def _where() -> str: - return socket.gethostname() - - -@env.task -async def quick() -> str: - """Placement probe: lands, sleeps 5s, reports where it ran (T0, T1, and after every --activate).""" - host = _where() - print(f"quick: running on {host}") - await asyncio.sleep(5) - return host - - -@env.task -async def sleep_for(seconds: int = 1800) -> str: - """Long runner for T1-T5 and T7. Heartbeats every 10s so a reset shows as a hostname change in the logs.""" - host = _where() - start = time.monotonic() - print(f"sleep_for: started on {host}, will run {seconds}s") - while (elapsed := time.monotonic() - start) < seconds: - await asyncio.sleep(min(10, seconds - elapsed)) - print(f"sleep_for: {host} alive {int(time.monotonic() - start)}s/{seconds}s") - print(f"sleep_for: finished on {host}") - return host - - -@env.task -async def fan_out(n: int = 50, seconds: int = 300) -> int: - """ - Parent with n concurrent sleep_for children in the same queue. - - T6: --n 50 --seconds 300, then force-drain and restart the owning leasor shard; - every child should end on attempt 2, never 3. - T7: --n 3 --seconds 1800 on the co-named queue; the parent and all children must - FAIL (not requeue) and the run must end, with no child left Unassigned. - """ - print(f"fan_out: parent on {_where()}, spawning {n} children x {seconds}s") - hosts = await asyncio.gather(*(sleep_for(seconds) for _ in range(n))) - by_host: dict[str, int] = {} - for h in hosts: - by_host[h] = by_host.get(h, 0) + 1 - print(f"fan_out: children finished on {by_host}") - return len(hosts) - - -if __name__ == "__main__": - flyte.init_from_config() - run = flyte.run(quick) - print(run.name, run.url) diff --git a/examples/cluster_drain/resource_exhausted.py b/examples/cluster_drain/resource_exhausted.py deleted file mode 100644 index fb4e817f5..000000000 --- a/examples/cluster_drain/resource_exhausted.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Exercise RESOURCE_EXHAUSTED handling against a depth-limited queue. - -Run with a queue configured with run_concurrency=1, action_concurrency=2, -and depth=5: - - flyte run --queue examples/cluster_drain/resource_exhausted.py main -""" - -import asyncio -import socket -from datetime import datetime, timezone - -import flyte - -env = flyte.TaskEnvironment( - name="resource-exhausted", - resources=flyte.Resources(cpu=1, memory="256Mi"), -) - - -@env.task -async def hold(i: int, seconds: int) -> int: - host = socket.gethostname() - started = datetime.now(timezone.utc).isoformat(timespec="seconds") - print(f"hold {i}: started on {host} at {started}", flush=True) - await asyncio.sleep(seconds) - finished = datetime.now(timezone.utc).isoformat(timespec="seconds") - print(f"hold {i}: finished on {host} at {finished}", flush=True) - return i - - -@env.task -async def main(count: int = 8, seconds: int = 5) -> list[int]: - tasks: list[asyncio.Task[int]] = [] - for i in range(count): - tasks.append(asyncio.create_task(hold(i, seconds))) - await asyncio.sleep(0.5) - results = await asyncio.gather(*tasks) - expected = list(range(count)) - if results != expected: - raise AssertionError(f"expected {expected}, got {results}") - print(f"all {count} actions completed", flush=True) - return results - - -if __name__ == "__main__": - flyte.init_from_config() - run = flyte.run(main, count=8, seconds=5) - print(run.name, run.url) - run.wait()