[serve] recursive cancellation #47873

zcin · 2024-09-30T23:44:48Z

Why are these changes needed?

Implement full recursive cancellation at Serve layer.

Currently, at the time that a parent request is cancelled, if:

The child request has not been assigned to a replica, meaning the replica scheduler in the router is still processing that request
AND the parent request is not directly blocking and waiting on the result of the child request

Then the child request will not be correctly cancelled. Note that if the second bullet point were actually flipped, i.e. the parent request was waiting on the result at the time of cancellation, then an asyncio.CancelledError will be sent directly to the replica scheduler and the child request will get successfully cancelled. However obviously this cannot be guaranteed to be true.

This PR adds a cache for asyncio tasks that are launched in the router to assign a request to a replica. When a parent request is cancelled, if there are any "pending assignment requests" in the cache, they are cancelled.

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: Cindy Zhang <[email protected]>

edoakes · 2024-10-25T20:14:00Z

python/ray/serve/tests/test_cancellation.py

@@ -365,5 +365,104 @@ async def get_out_of_band_response(self):
    assert h.get_out_of_band_response.remote().result() == "ok"


+def test_recursive_cancellation_during_execution(serve_instance):


in order for these tests to be meaningful, don't we need to turn off the core cancellation support?

So at first I took the approach of implementing full recursive cancellation at the Serve layer, but that leads to more of a performance drop. I think it's unnecessary to cover the part that core already covers (recursive task cancellation). We just need to take care of canceling the asyncio request assignment task before a request has been assigned to a replica.

So in this test file test_recursive_cancellation_during_execution already passes, and test_recursive_cancellation_during_assignment is the one that doesn't pass without this PR.

python/ray/serve/_private/replica.py

GeneDer · 2024-10-26T22:16:58Z

python/ray/serve/_private/router.py

@@ -594,6 +599,17 @@ async def assign_request(
    ) -> ReplicaResult:
        """Assign a request to a replica and return the resulting object_ref."""

+        response_id = uuid.uuid4()


Non-blocker: There is this generate_request_id() maybe we can rename and reuse here?

python/ray/serve/context.py

GeneDer · 2024-10-26T22:23:21Z

python/ray/serve/_private/router.py

+        response_id = uuid.uuid4()
+        assign_request_task = asyncio.current_task()
+        ray.serve.context._add_request_pending_assignment(
+            request_meta.request_id, response_id, assign_request_task


We should probably use internal_request_id for this. request_id can be set by the client from the header and can be repeated/ other unexpected behavior.

nice thanks for the suggestion, changed to use internal request id!

python/ray/serve/_private/router.py

python/ray/serve/context.py

Signed-off-by: Cindy Zhang <[email protected]>

GeneDer

Signed-off-by: Cindy Zhang <[email protected]>

GeneDer

Thanks Cindy!

Signed-off-by: Cindy Zhang <[email protected]>

## Why are these changes needed? Implement full recursive cancellation at Serve layer. Currently, at the time that a parent request is cancelled, if: - The child request has not been assigned to a replica, _meaning_ the replica scheduler in the router is still processing that request - AND the parent request is not directly blocking and waiting on the result of the child request Then the child request will not be correctly cancelled. Note that if the second bullet point were actually flipped, i.e. the parent request _was_ waiting on the result at the time of cancellation, then an `asyncio.CancelledError` will be sent directly to the replica scheduler and the child request will get successfully cancelled. However obviously this cannot be guaranteed to be true. This PR adds a cache for asyncio tasks that are launched in the router to assign a request to a replica. When a parent request is cancelled, if there are any "pending assignment requests" in the cache, they are cancelled. Signed-off-by: Cindy Zhang <[email protected]>

## Why are these changes needed? Implement full recursive cancellation at Serve layer. Currently, at the time that a parent request is cancelled, if: - The child request has not been assigned to a replica, _meaning_ the replica scheduler in the router is still processing that request - AND the parent request is not directly blocking and waiting on the result of the child request Then the child request will not be correctly cancelled. Note that if the second bullet point were actually flipped, i.e. the parent request _was_ waiting on the result at the time of cancellation, then an `asyncio.CancelledError` will be sent directly to the replica scheduler and the child request will get successfully cancelled. However obviously this cannot be guaranteed to be true. This PR adds a cache for asyncio tasks that are launched in the router to assign a request to a replica. When a parent request is cancelled, if there are any "pending assignment requests" in the cache, they are cancelled. Signed-off-by: Cindy Zhang <[email protected]> Signed-off-by: mohitjain2504 <[email protected]>

zcin added 14 commits September 26, 2024 14:57

fix recursive cancelation in serve

679a91d

Signed-off-by: Cindy Zhang <[email protected]>

Merge branch 'master' into fix-recursive-cancel

dcd60ff

Signed-off-by: Cindy Zhang <[email protected]>

Merge branch 'master' into fix-recursive-cancel

6cef130

Signed-off-by: Cindy Zhang <[email protected]>

lint

fe7bd70

Signed-off-by: Cindy Zhang <[email protected]>

up

ce4042b

Signed-off-by: Cindy Zhang <[email protected]>

remove on cancel

ebfaba3

Signed-off-by: Cindy Zhang <[email protected]>

Merge branch 'master' into fix-recursive-cancel

8ac7d12

Signed-off-by: Cindy Zhang <[email protected]>

fixes

e59e081

Signed-off-by: Cindy Zhang <[email protected]>

Merge branch 'master' into fix-recursive-cancel

ee9537f

Signed-off-by: Cindy Zhang <[email protected]>

fix lint + init typo

6d20408

Signed-off-by: Cindy Zhang <[email protected]>

remove unused code

7b5e6b5

Signed-off-by: Cindy Zhang <[email protected]>

Merge branch 'master' into fix-recursive-cancel

f95c889

Signed-off-by: Cindy Zhang <[email protected]>

only cancel assignment task

24a89bf

Signed-off-by: Cindy Zhang <[email protected]>

Merge branch 'master' into fix-recursive-cancel

4487c78

Signed-off-by: Cindy Zhang <[email protected]>

zcin marked this pull request as ready for review October 25, 2024 19:05

zcin requested review from GeneDer and edoakes October 25, 2024 19:06

lint

ed16fa7

Signed-off-by: Cindy Zhang <[email protected]>

edoakes reviewed Oct 25, 2024

View reviewed changes

GeneDer reviewed Oct 26, 2024

View reviewed changes

zcin added 2 commits October 28, 2024 11:46

address comments

3d7bcf1

Signed-off-by: Cindy Zhang <[email protected]>

Merge branch 'master' into fix-recursive-cancel

708c7de

Signed-off-by: Cindy Zhang <[email protected]>

edoakes approved these changes Oct 28, 2024

View reviewed changes

GeneDer approved these changes Oct 28, 2024

View reviewed changes

fix

3590a75

Signed-off-by: Cindy Zhang <[email protected]>

zcin added the go add ONLY when ready to merge, run all tests label Oct 28, 2024

GeneDer approved these changes Oct 28, 2024

View reviewed changes

fix unit test

c36c082

Signed-off-by: Cindy Zhang <[email protected]>

zcin merged commit cd4c34a into ray-project:master Oct 29, 2024
5 checks passed

zcin mentioned this pull request Oct 29, 2024

[serve] recursive cancellation #47935

Closed

zcin mentioned this pull request Oct 29, 2024

[Serve] Look into request got stuck issue #47585

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[serve] recursive cancellation #47873

[serve] recursive cancellation #47873

zcin commented Sep 30, 2024 •

edited

Loading

edoakes Oct 25, 2024

zcin Oct 25, 2024

GeneDer Oct 26, 2024

GeneDer Oct 26, 2024

zcin Oct 28, 2024

GeneDer left a comment

GeneDer left a comment

		@@ -365,5 +365,104 @@ async def get_out_of_band_response(self):
		assert h.get_out_of_band_response.remote().result() == "ok"


		def test_recursive_cancellation_during_execution(serve_instance):

[serve] recursive cancellation #47873

[serve] recursive cancellation #47873

Conversation

zcin commented Sep 30, 2024 • edited Loading

Why are these changes needed?

Related issue number

Checks

edoakes Oct 25, 2024

Choose a reason for hiding this comment

zcin Oct 25, 2024

Choose a reason for hiding this comment

GeneDer Oct 26, 2024

Choose a reason for hiding this comment

GeneDer Oct 26, 2024

Choose a reason for hiding this comment

zcin Oct 28, 2024

Choose a reason for hiding this comment

GeneDer left a comment

Choose a reason for hiding this comment

GeneDer left a comment

Choose a reason for hiding this comment

zcin commented Sep 30, 2024 •

edited

Loading