[pardeep-singh]

Description

RequestRouter._probe_queue_lens() cancels every get_queue_len probe in its finally block but returns without waiting for them. Cancellation is cooperative, so a probe that handles CancelledError with async cleanup is still running after the router has dropped its result and gone on to fulfil the request or start the next attempt.

This change:

• Keeps master's cancel-all loop. The finally still cancels every probe. cancel() on a probe that already raised also keeps its exception from being logged as "Task exception was never retrieved". Unfinished probes also get a done callback that logs (and so retrieves) an exception raised while they handle the cancellation.
• Processes the results first. The pending and done probes are handled as before: timed-out probes are reported as unavailable, and ActorDiedError / ActorUnavailableError update the replica state.
• Then waits for the cancelled probes with asyncio.wait, just before returning. The wait is bounded by the round's own queue_len_response_deadline_s, so a round takes at most twice its deadline. RunningReplica.get_queue_len requests cancellation of its RPC and re-raises CancelledError without awaiting the replica, so in production the wait ends within one event loop turn. A probe that ignores cancellation for longer is logged and left behind, so it can't block routing.
• No wait on caller cancellation. If the routing task itself is cancelled (e.g. its request was fulfilled by another task), it exits right away. A cancellation during the bounded wait doesn't drop results that were already processed.

Related issues

Fixes #65719. The earlier PR #65720 was closed as stale; it waited only on the deadline path and without a bound.

Additional information

Tests in test_pow_2_request_router.py, using a fake replica whose get_queue_len does async cleanup when cancelled:

• test_cancelled_get_queue_len_finishes_before_routing: the slow-to-cancel replica's cleanup has finished by the time routing returns the other replica.
• test_slow_to_cancel_get_queue_len_does_not_block_routing: a probe whose cleanup takes far longer than the bound doesn't delay routing past the bound: with a 100s max deadline, routing still finishes in under 1s.
• test_routing_task_cancelled_while_cancelled_probe_finishes: cancelling the request while its probe is cleaning up still raises CancelledError, and the routing task ends cancelled in well under the wait's bound.
• test_failed_probe_exception_retrieved_when_routing_task_cancelled: a probe that has already raised when the routing task is cancelled isn't reported as "never retrieved".
• test_probe_results_kept_when_cancelled_during_probe_cleanup: cancelling during the bounded wait exits in well under the bound and still evicts the replica whose probe returned ActorDiedError.
• test_cancelled_probe_error_is_logged: an exception a cancelled probe raises while finishing is logged as a warning and not reported as "never retrieved". The test waits for the warning itself rather than relying on the router's wait.
• test_cancelled_routing_task_does_not_wait_for_probe_cleanup: cancellation returns promptly while a cancelled probe is slow to clean up.

The last three were added after review and each fails on the earlier version of this PR. The two cancel-during-wait tests fail if the wait holds a cancellation until its bound.

In test_actor_replica_wrapper.py, test_get_queue_len_cancellation_cancels_rpc cancels RunningReplica.get_queue_len against a real actor. The actor holds its cancellation open until the test releases it, so the test checks that the probe is cancelled after one event loop turn, without awaiting the replica, and that the cancellation request reaches the replica.

pytest python/ray/serve/tests/unit/test_pow_2_request_router.py python/ray/serve/tests/unit/test_consistent_hash_router.py   # 168 passed
pytest python/ray/serve/tests/test_actor_replica_wrapper.py   # 17 passed (local CPU cluster)
pre-commit run --files <changed files>   # all hooks pass

Why this isn't a duplicate: the earlier PR #65720 was closed by the stale bot; no other open PR references #65719.

AI assistance: I used an AI coding assistant to write this change, its tests and this description.

[gemini-code-assist]

Code Review

This pull request improves task cancellation handling in _probe_queue_lens by ensuring unfinished probe tasks are cancelled, their exceptions are retrieved to prevent "never retrieved" logs, and the router waits for them to finish up to a maximum deadline. It also adds comprehensive unit tests to verify these cancellation behaviors. The feedback suggests improving the readability of the done callback lambda by using an explicit conditional expression instead of the or operator.

[johntaylor-cell]

Cancelling only the unfinished probes means a probe that has already failed when the routing task is cancelled (its request was cancelled or fulfilled by another task) never has its exception retrieved, so asyncio logs "Task exception was never retrieved". Let's keep master's loop that cancels every task, since cancel() on a finished task also suppresses that log, and add a test where one probe has already raised when the routing task is cancelled.

The new wait also runs before done is processed, so a cancel that lands during it drops results already in hand, e.g. an ActorDiedError from that round doesn't remove the replica. Processing the results first and doing the bounded wait just before return avoids that, and lets a cancelled routing task exit right away instead of waiting up to the max response deadline.

[pardeep-singh]

Thanks, both points were right. Fixed in fc45510:

• The finally is back to master's loop that cancels every probe, so a probe that already raised has its exception retrieved. Unfinished probes still get the done-callback that retrieves an exception raised while handling the cancellation.
• The bounded wait for the cancelled probes now runs after pending and done are processed, just before return. A cancellation during the wait no longer drops results such as an ActorDiedError, and a routing task cancelled while its probes are in flight exits right away.

New tests, all three failing on the previous commit:

• test_failed_probe_exception_retrieved_when_routing_task_cancelled: one probe has already raised when the routing task is cancelled; no "never retrieved" is reported.
• test_probe_results_kept_when_cancelled_during_probe_cleanup: cancelling during the bounded wait still evicts the replica whose probe returned ActorDiedError.
• test_cancelled_routing_task_does_not_wait_for_probe_cleanup: cancellation returns promptly while a cancelled probe is slow to clean up.

test_pow_2_request_router.py and test_consistent_hash_router.py: 167 passed (3 runs).

[johntaylor-cell]

Thanks, both points are fixed. A few more comments:

1. RunningReplica.get_queue_len, the only production probe, already cancels synchronously, so the wait only adds latency; consider documenting that probes must cancel promptly and dropping it.
2. If the wait stays, bound it by the round's queue_len_response_deadline_s rather than the max, which today stretches a 0.1s round to 1.1s.
3. The done callback silently drops exceptions raised during cancellation; please log them and add a test, since removing it passes every test today.
4. test_slow_to_cancel_get_queue_len_does_not_block_routing allows 10s against a 0.01s bound; assert the elapsed time instead.

[pardeep-singh]

@johntaylor-cell thanks for the second pass. Pushed e8c0060:

1. I kept the wait, since [Serve] Drain cancelled queue-length probes after response deadlines #65719 asks for cancelled probes to finish before the router moves on, and documented that probes should handle cancellation promptly. RunningReplica.get_queue_len cancels its RPC synchronously, so in production the wait ends within one event loop turn. If you or a Serve maintainer would rather drop it and rely on the docs alone, I can do that.
2. The wait is now bounded by the round's queue_len_response_deadline_s, not the max, so a 0.1s round stays a 0.1s round.
3. The done callback now logs an exception a cancelled probe raises while finishing (_log_cancelled_probe_error). The new test test_cancelled_probe_error_is_logged checks the warning and that nothing is reported as "never retrieved". It fails without the callback.
4. test_slow_to_cancel_get_queue_len_does_not_block_routing sets the max deadline to 100s and asserts that routing finishes in under 1s. It fails against the previous version, which waited for the max.

test_pow_2_request_router.py: 141 passed. Could you take another look when you have a moment?

[johntaylor-cell]

@pardeep-singh Thanks, 2, 3 and 4 look good. A few more should do it:

1. The cancel-promptly rule only lives in a comment inside _probe_queue_lens. Please put it on the RunningReplica.get_queue_len docstring, since that's the public probe interface, and add a test of its cancel path.
2. With the bound at 1s, the two cancel-during-wait tests pass even if the wait holds the cancel until its bound. Please assert the task exits well under the bound after cancel().
3. wait_for(task, timeout=100) outlasts the file's 60s Bazel timeout, so a regression would show as a target timeout. A few seconds should be enough.
4. test_cancelled_probe_error_is_logged only sees the warning because the trailing wait runs the callback first. Please wait for the warning inside the patch block.

[pardeep-singh]

@johntaylor-cell thanks, all four are addressed in 9c68c98:

1. The cancel-promptly rule is now on the RunningReplica.get_queue_len docstring. The new test_get_queue_len_cancellation_cancels_rpc in test_actor_replica_wrapper.py cancels it while the actor is inside get_num_ongoing_requests. It checks the cancellation finishes in under 1s and the call is cancelled on the replica.
2. Both cancel-during-wait tests now use a 2s bound and assert the routing task exits in under 0.5s after cancel(). In test_routing_task_cancelled_while_cancelled_probe_finishes the timing covers the routing task, waiting for curr_num_routing_tasks == 0, since the _choose_replica_for_request caller returns before the routing task does. I checked both by making the wait sleep until its bound on cancellation: both fail.
3. test_slow_to_cancel_get_queue_len_does_not_block_routing now uses wait_for(task, timeout=5).
4. test_cancelled_probe_error_is_logged now waits for the warning inside the patch block.

test_pow_2_request_router.py + test_consistent_hash_router.py: 168 passed. The description is updated to match.

[johntaylor-cell]

We are getting close @pardeep-singh. A few small ones remain.

1. The new get_queue_len docstring says the router waits for the probe even when the routing task is cancelled, and before acting on the results, but it does neither. Something like "The router cancels this at the response deadline or with its routing task; re-raise CancelledError without awaiting the replica" is enough.
2. test_get_queue_len_cancellation_cancels_rpc still passes if the handler awaits the cancelled ref, since the fake actor finishes its cancellation after one RPC. Please have it hold the cancellation open until the test releases it.
3. test_routing_task_cancelled_while_cancelled_probe_finishes passes even if the routing task swallows its cancellation; please assert it ends cancelled.
4. The 2s deadlines add about 3s to the file without catching more. 1s already fails the hold-until-bound case.

[pardeep-singh]

Thanks @johntaylor-cell, all four are in 0352dfb:

1. The get_queue_len docstring now says the router cancels it at the response deadline or with its routing task, and that it re-raises CancelledError without awaiting the replica.
2. BlockingQueueLenActor holds its cancellation open until the test sends a release signal. With get_queue_len changed to await the cancelled ref, the test now times out and fails.
3. test_routing_task_cancelled_while_cancelled_probe_finishes grabs the routing task and asserts routing_task.cancelled(). With the bounded wait changed to swallow CancelledError, it fails.
4. Both cancel-during-wait tests use a 1s deadline.

Locally: the router unit files pass (168), and test_actor_replica_wrapper.py passes (17).

[johntaylor-cell]

Thanks @pardeep-singh, all four look good. Two small ones and this looks ready to me:

1. test_get_queue_len_cancellation_cancels_rpc still passes if the handler awaits the cancelled ref with a timeout under 5s. await asyncio.sleep(0) then assert task.done() right after task.cancel() pins it without any timing.
2. ray.cancel only drops a queued call and can't stop the replica's sync get_num_ongoing_requests once it runs, so "the RPC is cancelled" in the docstring and test should say cancellation of the RPC is requested.

[pardeep-singh]

Thanks @johntaylor-cell, both done in 6b2aa95:

1. test_get_queue_len_cancellation_cancels_rpc now does await asyncio.sleep(0) and assert task.cancelled() right after task.cancel(), with no timeout. I checked it fails if the handler awaits the cancelled ref with a 1s timeout before re-raising. It still waits for the replica's cancelled signal afterwards.
2. The get_queue_len docstring, the test docstring and comment, and the PR description now say cancellation of the RPC is requested.

pytest test_actor_replica_wrapper.py -k 'get_queue_len_cancellation or reserve_slot_cancellation': 2 passed; pre-commit passes.

[johntaylor-cell]

Thanks @pardeep-singh, both look good. One last nit is that the comment in _probe_queue_lens still says RunningReplica.get_queue_len "cancels its RPC synchronously", and the test is still named ..._cancels_rpc; worth the same "requests cancellation" wording. And please fix the premerge CI failure if necessary. Looks good to me otherwise.