[Serve] Revisiting `ProxyState` to fix draining sequence #41722

alexeykudinkin · 2023-12-08T04:55:58Z

Why are these changes needed?

These changes are revisiting ProxyState to

Address an issue w/ a draining sequence ([Serve] Draining Proxy is shutdown prematurely #41726), when Proxy is shutdown prematurely
Simplify ProxyState reconciliation sequence

Changes

Offloaded async interactions w/ the ProxyActor from ProxyState to ActorProxyWrapper (substantially simplifying it)
Streamlined overall reconciliation sequence to make it easier to comprehend
Cleaned up tests (avoid sleeping, added validations), made sure all flows are adequately covered

Related issue number

Closes #41726 (added test)

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 :(

edoakes

It's very difficult to see where the actual behavior fix is here -- let's split out the targeted bug fix and the general refactoring into two separate PRs please. This is currently a very large change that adds unnecessary risk (especially if cherry-picked).

(Let me know if I'm mistaken and all of the refactoring is necessary to fix the issue)

edoakes · 2023-12-08T17:13:59Z

python/ray/serve/tests/test_proxy_actor_wrapper.py

+    # Yield the event-loop
+    await asyncio.sleep(0.001)
+
+    assert aio_fut.done()
+    assert aio_fut.result() == "test"


You can use ray._private.test_utils.async_wait_for_condition as a common pattern for adding these assertions that require the loop to be yielded

python/ray/serve/_private/proxy_state.py

GeneDer

Thanks for revisit the code and trying to make it better! I think the biggest point is would be nice to keep ProxyWrapperCallStatus around and revert the update() method name. The logics seems to be consistent as before, but hopefully any regression will be caught by the CI (currently passing 😄)

python/ray/serve/_private/constants.py

python/ray/serve/_private/proxy_state.py

GeneDer · 2023-12-08T23:09:57Z

python/ray/serve/_private/proxy_state.py

+            self._ready_check_future = wrap_as_future(
+                self._actor_handle.ready.remote(), timeout_s=timeout_s
+            )
+            return None


Nit (non-blocker): probably don't need to explicitly return None here. We can directly go into checking whether the future is done below. If the ready call is fast enough, this will help save a loop and start time a bit.

Good call! Will update

Just to call out that it won't really pay off since this code executing is occupying the event-loop, so it only will save us a loop when we create a future that immediately returns (ie failed). But i like general direction of the thinking here and i don't think i see the downside of doing it

python/ray/serve/_private/proxy_state.py

GeneDer · 2023-12-09T00:05:11Z

python/ray/serve/_private/proxy_state.py

+            return None
+
+    def is_drained(
+        self, timeout_s: float = PROXY_HEALTH_CHECK_TIMEOUT_S


This should probably be PROXY_DRAIN_CHECK_PERIOD_S XP

Or make a timeout for drain check?

We can generalize HC timeout (i don't think we need a separate one for everything, provided they serve practically identical purpose)

I think if we decided to just want to use PROXY_HEALTH_CHECK_TIMEOUT_S for all the timeouts, maybe we should add a comment to that constant or preferably rename the constant.

Also is there a reason we should add a default for is_drained but not for is_ready and is_healthy?

Yeah, good call. Will make them coherent.

python/ray/serve/_private/proxy_state.py

#41744) (#41755) This is a minified version of https://github.com/ray-project/ray/pull/41722/files, specifically put to be cherry-picked into 2.9 Addresses #41726 Addresses following gaps: Patches ActorProxyWrapper.is_drained method to handle RPC response properly Cherry-picks a test from [Serve] Revisiting ProxyState to fix draining sequence #41722 to validate draining sequence is correct --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin · 2024-01-10T22:07:20Z

@edoakes @GeneDer updated/addressed comments, PTAL

GeneDer · 2024-01-10T23:27:35Z

python/ray/serve/_private/proxy_state.py

+            return None
+
+    def is_drained(
+        self, timeout_s: float = PROXY_HEALTH_CHECK_TIMEOUT_S


I think if we decided to just want to use PROXY_HEALTH_CHECK_TIMEOUT_S for all the timeouts, maybe we should add a comment to that constant or preferably rename the constant.

Also is there a reason we should add a default for is_drained but not for is_ready and is_healthy?

python/ray/serve/_private/proxy_state.py

edoakes · 2024-01-11T23:02:33Z

Is this ready for review @alexeykudinkin ?

alexeykudinkin · 2024-01-11T23:15:45Z

Is this ready for review @alexeykudinkin ?

Yes

GeneDer

I still feel we should keep ProxyWrapperCallStatus enum around, but not a hard blocker. Thanks for revising this logic! :)

python/ray/serve/_private/proxy_state.py

edoakes

Looks great! Just a few stylistic comments and suggestions aroundwrap_as_future.

python/ray/serve/_private/proxy_state.py

edoakes · 2024-01-12T22:50:45Z

python/ray/serve/_private/proxy_state.py

+        if self._ready_check_future is None:
+            self._ready_check_future = wrap_as_future(
+                self._actor_handle.ready.remote(), timeout_s=timeout_s
+            )


maybe an unnecessary optimization, but should we just early-return None here?

Ha! That's actually what it's been before, but @GeneDer made a good point on just letting it fall through so that if future has failed right away we don't wait for another iteration

https://github.com/ray-project/ray/pull/41722/files#r1421100067

I think it's this thread #41722 (comment) But I don't have strong preference for either 😛

python/ray/serve/_private/proxy_state.py

edoakes · 2024-01-12T23:01:20Z

python/ray/serve/_private/proxy_state.py

+        TimeoutError(f"Future cancelled after timeout {timeout_s}s"),
+    )
+    # Cancel timeout handler upon completion of the future
+    aio_fut.add_done_callback(lambda _: timeout_handler.cancel())


I am pretty sure that because the timeout_handler does not catch TimeoutError, it can log a nasty exception traceback ("unhandled exception in coroutine ..."). Can you add an except TimeoutError: pass block to the timeout handler?

timeout_handler is just a handler for _try_set_exception (_try_cancel_future previously) why would it need to handle TimeoutError?

Oh I'm sorry, typo, I meant asyncio.CancelledError

Got it. This is actually a TimerHandle not the Future, so cancelling it actually won't be throwing CancelledError:

def cancel(self): if not self._cancelled: self._cancelled = True if self._loop.get_debug(): # Keep a representation in debug mode to keep callback and # parameters. For example, to log the warning # "Executing <Handle...> took 2.5 second" self._repr = repr(self) self._callback = None self._args = None

alexeykudinkin · 2024-01-18T06:06:37Z

@edoakes feedback addressed, PTAL

edoakes · 2024-01-19T15:59:03Z

@alexeykudinkin DCO and test failure

GeneDer · 2024-01-19T17:23:01Z

@alexeykudinkin just FYI that test failure has to do with the latest Starlette version. If you just rebase to the latest master, it should automatically skipping that test and allow merge. We have a separate PR to fix the starlette issue #42378

alexeykudinkin · 2024-01-20T00:05:01Z

@edoakes @GeneDer addressed feedback, PTAL

GeneDer · 2024-01-20T00:16:21Z

@alexeykudinkin I think those are already approved by us right? Is there anything new that you want us to look in particular? DCO is still failing bc the email you used in that commit is alexey.kudinkin@gmail.com instead of ak@anyscale.com. I think you can amend that commit with the anyscale email and force push again then it should be resolved :)

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Rebased `set_status` calls to invoke `try_update_status` instead (for proper state-management); Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

…operly drained; Missing changes Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Clean up `ProxyWrapperCallStatus` Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Added logging Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com> Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin · 2024-01-20T06:42:52Z

Sorry, i missed that @edoakes had already approved it as well.

Thanks for reviewing it!

alexeykudinkin force-pushed the ak/srv-drn-fix branch from e0080b8 to 73909e2 Compare December 8, 2023 08:16

alexeykudinkin changed the title ~~[WIP][Serve] Revisiting ProxyState to fix draining sequence, simplify state handling~~ [WIP][Serve] Revisiting ProxyState to fix draining sequence Dec 8, 2023

alexeykudinkin linked an issue Dec 8, 2023 that may be closed by this pull request

[Serve] Draining Proxy is shutdown prematurely #41726

Closed

alexeykudinkin force-pushed the ak/srv-drn-fix branch from 73909e2 to d924379 Compare December 8, 2023 08:24

alexeykudinkin requested review from GeneDer and edoakes December 8, 2023 08:24

edoakes requested changes Dec 8, 2023

View reviewed changes

alexeykudinkin changed the title ~~[WIP][Serve] Revisiting ProxyState to fix draining sequence~~ [Serve] Revisiting ProxyState to fix draining sequence Dec 8, 2023

This was referenced Dec 8, 2023

[Serve] Patching ActorProxyWrapper to properly handle is_drained RPC #41744

Merged

[Serve] Patching ActorProxyWrapper to properly handle is_drained RPC #41755

Merged

GeneDer reviewed Dec 9, 2023

View reviewed changes

alexeykudinkin force-pushed the ak/srv-drn-fix branch from 8faf7f7 to 74dcaf4 Compare January 10, 2024 18:43

GeneDer reviewed Jan 11, 2024

View reviewed changes

GeneDer approved these changes Jan 12, 2024

View reviewed changes

python/ray/serve/_private/proxy_state.py Outdated Show resolved Hide resolved

edoakes reviewed Jan 12, 2024

View reviewed changes

edoakes approved these changes Jan 19, 2024

View reviewed changes

alexeykudinkin force-pushed the ak/srv-drn-fix branch from b700fba to b947610 Compare January 19, 2024 17:56

alexeykudinkin added 4 commits January 19, 2024 22:41

Merge start_new_ready_check into is_ready

39c1447

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Cleaned up try_update_status;

0231ccf

Rebased `set_status` calls to invoke `try_update_status` instead (for proper state-management); Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Streamlined is_healthy, is_ready checks in ProxyActorWrapper

0ebfe1b

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

b6b3660

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin and others added 25 commits January 19, 2024 22:41

Simplified is_healthy, is_ready checks even more

feffe80

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed is_drained check to properly return whether proxy actor is pr…

2144898

…operly drained; Missing changes Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Streamlined proxy's state reconciliation flow

d04d010

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Simplified readiness, health, drain checks;

ab1cc40

Clean up `ProxyWrapperCallStatus` Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Make sure exceptions are properly handled in proxy state reconciliation

5ba6068

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Cleaned up docs

d9ef34b

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Added test for ActorProxyWrapper

e7c5e1e

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed typos;

759d904

Added logging Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed tests

953343a

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed more tests

73d881e

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed some more tests

9e05832

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

8d0cdcd

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixing typo

136df1c

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

After rebase clean-up

189cad4

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

c8dcb57

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Make proxy health-check configurable

06b1755

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Unify is_drained API to require timeout be explicitly provided

f4283df

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Missing log statement

b787188

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

adfd2f3

Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com> Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Cleaned up exception handling

eff6f23

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

_try_cancel_future -> _try_set_exception

959feba

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Make timeout_s optional

c5d2034

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixing typo

7002e11

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixing another typo

cc99b4a

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixing tests

8a1e954

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin force-pushed the ak/srv-drn-fix branch from b947610 to 8a1e954 Compare January 20, 2024 06:42

edoakes merged commit e7ae6d0 into master Jan 22, 2024
2 checks passed

edoakes deleted the ak/srv-drn-fix branch January 22, 2024 01:26

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Serve] Revisiting `ProxyState` to fix draining sequence #41722

[Serve] Revisiting `ProxyState` to fix draining sequence #41722

alexeykudinkin commented Dec 8, 2023 •

edited by edoakes

Loading

edoakes left a comment

edoakes Dec 8, 2023

GeneDer left a comment

GeneDer Dec 8, 2023

alexeykudinkin Dec 11, 2023

alexeykudinkin Jan 10, 2024

GeneDer Dec 9, 2023

GeneDer Dec 9, 2023

alexeykudinkin Dec 11, 2023 •

edited

Loading

GeneDer Jan 10, 2024

alexeykudinkin Jan 11, 2024

alexeykudinkin commented Jan 10, 2024

GeneDer Jan 10, 2024

edoakes commented Jan 11, 2024

alexeykudinkin commented Jan 11, 2024

GeneDer left a comment

edoakes left a comment

edoakes Jan 12, 2024

alexeykudinkin Jan 16, 2024

GeneDer Jan 16, 2024

edoakes Jan 12, 2024

alexeykudinkin Jan 18, 2024

edoakes Jan 18, 2024

alexeykudinkin Jan 18, 2024

alexeykudinkin commented Jan 18, 2024

edoakes commented Jan 19, 2024

GeneDer commented Jan 19, 2024 •

edited

Loading

alexeykudinkin commented Jan 20, 2024

GeneDer commented Jan 20, 2024

alexeykudinkin commented Jan 20, 2024

[Serve] Revisiting ProxyState to fix draining sequence #41722

[Serve] Revisiting ProxyState to fix draining sequence #41722

Conversation

alexeykudinkin commented Dec 8, 2023 • edited by edoakes Loading

Why are these changes needed?

Related issue number

Checks

edoakes left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

GeneDer left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alexeykudinkin Dec 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alexeykudinkin commented Jan 10, 2024

Choose a reason for hiding this comment

edoakes commented Jan 11, 2024

alexeykudinkin commented Jan 11, 2024

GeneDer left a comment

Choose a reason for hiding this comment

edoakes left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alexeykudinkin commented Jan 18, 2024

edoakes commented Jan 19, 2024

GeneDer commented Jan 19, 2024 • edited Loading

alexeykudinkin commented Jan 20, 2024

GeneDer commented Jan 20, 2024

alexeykudinkin commented Jan 20, 2024

[Serve] Revisiting `ProxyState` to fix draining sequence #41722

[Serve] Revisiting `ProxyState` to fix draining sequence #41722

alexeykudinkin commented Dec 8, 2023 •

edited by edoakes

Loading

alexeykudinkin Dec 11, 2023 •

edited

Loading

GeneDer commented Jan 19, 2024 •

edited

Loading