[Serve] Patching ActorProxyWrapper to properly handle is_drained RPC (

ray-project#41744) 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 ray-project#41726 --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com>
ShuN6211 · Dec 8, 2023 · 567e574 · 567e574
1 parent d6d2689
commit 567e574
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 4 deletions.
diff --git a/python/ray/serve/_private/proxy_state.py b/python/ray/serve/_private/proxy_state.py
@@ -284,8 +284,15 @@ def is_drained(self) -> ProxyWrapperCallStatus:
         finished, _ = ray.wait([self._is_drained_obj_ref], timeout=0)
         if finished:
             self._is_drained_obj_ref = None
-            ray.get(finished[0])
-            return ProxyWrapperCallStatus.FINISHED_SUCCEED
+            is_drained = ray.get(finished[0])
+            if is_drained:
+                return ProxyWrapperCallStatus.FINISHED_SUCCEED
+            else:
+                # NOTE: Even though call returned successfully, we have to
+                #       report it as FINISHED_FAILED to make sure that
+                #       draining process doesn't move forward until draining
+                #       completes
+                return ProxyWrapperCallStatus.FINISHED_FAILED
         else:
             return ProxyWrapperCallStatus.PENDING
 

diff --git a/python/ray/serve/tests/test_proxy_state.py b/python/ray/serve/tests/test_proxy_state.py
@@ -53,12 +53,14 @@ def __init__(self, *args, **kwargs):
         self.actor_handle = FakeProxyActor(*args, **kwargs)
         self.ready = ProxyWrapperCallStatus.FINISHED_SUCCEED
         self.health = ProxyWrapperCallStatus.FINISHED_SUCCEED
+        self.drained = ProxyWrapperCallStatus.FINISHED_SUCCEED
         self.worker_id = "mock_worker_id"
         self.log_file_path = "mock_log_file_path"
         self.health_check_ongoing = False
         self.is_draining = False
         self.shutdown = False
         self.num_health_checks = 0
+        self.num_drain_checks = 0
 
     @property
     def actor_id(self) -> str:
@@ -74,7 +76,7 @@ def start_new_health_check(self):
         self.health_check_ongoing = True
 
     def start_new_drained_check(self):
-        pass
+        self.is_draining = True
 
     def is_ready(self) -> ProxyWrapperCallStatus:
         return self.ready
@@ -85,7 +87,9 @@ def is_healthy(self) -> ProxyWrapperCallStatus:
         return self.health
 
     def is_drained(self) -> ProxyWrapperCallStatus:
-        pass
+        self.num_drain_checks += 1
+        self.is_draining = False
+        return self.drained
 
     def is_shutdown(self):
         return self.shutdown
@@ -99,6 +103,9 @@ def kill(self):
     def get_num_health_checks(self):
         return self.num_health_checks
 
+    def get_num_drain_checks(self):
+        return self.num_health_checks
+
 
 def _create_proxy_state_manager(
     http_options: HTTPOptions = HTTPOptions(),
@@ -802,6 +809,52 @@ def check_worker_node_proxy_actor_is_removed():
     assert manager._proxy_states[HEAD_NODE_ID].status == ProxyStatus.HEALTHY
 
 
+@patch("ray.serve._private.proxy_state.PROXY_HEALTH_CHECK_PERIOD_S", 5)
+def test_proxy_state_reconcile_draining_success():
+    """Test that the proxy will remain DRAINING even if health check succeeds."""
+    timer = MockTimer(start_time=0)
+    # Start with HEALTHY state
+    proxy_state = _create_proxy_state(status=ProxyStatus.HEALTHY, timer=timer)
+    # Simulate health-checks passing
+    proxy_state._actor_proxy_wrapper.healthy = ProxyWrapperCallStatus.FINISHED_SUCCEED
+    # Simulate is_drained check returning false
+    proxy_state._actor_proxy_wrapper.drained = ProxyWrapperCallStatus.FINISHED_FAILED
+
+    for _ in range(10):
+        proxy_state.update(draining=True)
+        assert proxy_state.status == ProxyStatus.DRAINING
+        # Advance timer by 5 (to trigger new health-check, drain-check)
+        timer.advance(5)
+
+    # assert proxy_state._actor_proxy_wrapper.get_num_health_checks() == 10
+    # assert proxy_state._actor_proxy_wrapper.get_num_drain_checks() == 9
+
+    # Make sure the status is still DRAINING
+    assert proxy_state.status == ProxyStatus.DRAINING
+
+    # Simulate is_drained request to ProxyActor pending (for 5 iterations)
+    proxy_state._actor_proxy_wrapper.drained = ProxyWrapperCallStatus.PENDING
+
+    for _ in range(5):
+        proxy_state.update(draining=True)
+        assert proxy_state.status == ProxyStatus.DRAINING
+        # Advance timer by 5 (to trigger new health-check, drain-check)
+        timer.advance(5)
+
+    # assert proxy_state._actor_proxy_wrapper.get_num_health_checks() == 15
+    # No new drain checks will occur, since there's a pending one (not completed yet)
+    # assert proxy_state._actor_proxy_wrapper.get_num_drain_checks() == 14
+
+    # Simulate draining completed
+    proxy_state._actor_proxy_wrapper.drained = ProxyWrapperCallStatus.FINISHED_SUCCEED
+    # Advance timer by 5 (to trigger new health-check, drain-check on next iteration)
+    timer.advance(5)
+
+    proxy_state.update(draining=True)
+    # State should transition to DRAINED
+    assert proxy_state.status == ProxyStatus.DRAINED
+
+
 def test_is_ready_for_shutdown(all_nodes):
     """Test `is_ready_for_shutdown()` returns True the correct state.