From 4f969b77a4851f3caebb185cf5c5a1472597c21c Mon Sep 17 00:00:00 2001 From: harshit-anyscale Date: Mon, 24 Aug 2026 16:07:00 +0530 Subject: [PATCH 1/5] [serve] Close direct-ingress listeners once the deregistration window ends Behind HAProxy, a draining replica today keeps accepting new requests for as long as they arrive. An HAProxy reload leaves the old process running with a frozen backend list (it cannot be updated, only awaited), so a stale worker can keep sending fresh requests to the draining replica for minutes. When the node then goes away, those late requests are severed mid-flight and the client sees an unretryable error. Fix: behind HAProxy, drain in two phases -- serve the deregistration window in full, then close the direct-ingress HTTP listener and wait for in-flight requests. Late arrivals are refused at connect time, which HAProxy retries on another replica (retry-on conn-failure + option redispatch). Without HAProxy nothing changes: there is no retrying party in front, so a refusal would reach the client. Signed-off-by: harshit-anyscale Co-Authored-By: Claude Fable 5 --- python/ray/serve/_private/replica.py | 57 ++++++++++- .../serve/tests/unit/test_replica_quiesce.py | 98 ++++++++++++++++++- 2 files changed, 151 insertions(+), 4 deletions(-) diff --git a/python/ray/serve/_private/replica.py b/python/ray/serve/_private/replica.py index 6c0194f52df3..b09802bc8703 100644 --- a/python/ray/serve/_private/replica.py +++ b/python/ray/serve/_private/replica.py @@ -1180,6 +1180,10 @@ def __init__( def max_ongoing_requests(self) -> int: return self._deployment_config.max_ongoing_requests + @property + def _is_direct_ingress(self) -> bool: + return self._ingress and RAY_SERVE_ENABLE_DIRECT_INGRESS + def get_num_ongoing_requests(self) -> int: return self._metrics_manager.get_num_ongoing_requests() + len( self._reserved_slots @@ -2027,6 +2031,49 @@ async def _drain_ongoing_requests(self, min_draining_period_s: float = 0.0): ) break + def _stop_accepting_direct_ingress(self) -> None: + """Close the direct-ingress HTTP listener; in-flight requests finish. + + Only the HTTP listener: the direct-ingress gRPC server (if any) keeps + accepting until the post-drain quiesce, as before. + """ + if self._direct_ingress_http_server is not None: + self._direct_ingress_http_server.should_exit = True + + async def _drain_behind_haproxy(self, min_draining_period_s: float) -> None: + """Two-phase drain for replicas fronted by HAProxy. + + Stay fully reachable for the deregistration window, then close the + HTTP listener and wait for in-flight requests. A stale HAProxy worker + (an old soft-stopping process with a frozen backend list) can keep + routing here for minutes after a reload; refusing it at connect time + makes it retry another replica (`retry-on conn-failure` + `option + redispatch`) instead of admitting a request that would be severed + when the replica exits. + + Only safe with a retrying proxy in front -- without HAProxy the + refusal would reach the client, so callers use the plain drain there. + """ + # Phase 1: remain fully reachable for the deregistration window. + if min_draining_period_s > 0: + logger.info( + f"Draining: staying reachable for {min_draining_period_s:.1f}s " + "so load balancers can deregister this replica.", + extra={"log_to_stderr": False}, + ) + await asyncio.sleep(min_draining_period_s) + + # Phase 2: refuse new connections; stale routers retry elsewhere. + logger.info( + "Draining: closing the direct ingress HTTP listener; late arrivals " + "will be refused so the caller can retry another replica.", + extra={"log_to_stderr": False}, + ) + self._stop_accepting_direct_ingress() + + # Phase 3: wait for in-flight requests (window already served). + await self._drain_ongoing_requests() + async def shutdown(self): try: self._user_callable_wrapper.stop_user_loop_watchdog() @@ -2064,10 +2111,16 @@ def remaining_grace_s() -> float: # deregister it; the drain also waits for in-flight requests. min_draining_period_s = ( RAY_SERVE_DIRECT_INGRESS_MIN_DRAINING_PERIOD_S - if RAY_SERVE_ENABLE_DIRECT_INGRESS and self._ingress + if self._is_direct_ingress else 0.0 ) - await self._drain_ongoing_requests(min_draining_period_s) + if self._is_direct_ingress and RAY_SERVE_ENABLE_HA_PROXY: + # HAProxy retries refused connections, so stop accepting once + # the deregistration window is over. + await self._drain_behind_haproxy(min_draining_period_s) + else: + # No retrying party in front; a refusal would reach the client. + await self._drain_ongoing_requests(min_draining_period_s) # Requests can still arrive after the drain (stale routers, keep-alive # connections). Quiesce before reporting shutdown complete: reject new diff --git a/python/ray/serve/tests/unit/test_replica_quiesce.py b/python/ray/serve/tests/unit/test_replica_quiesce.py index 331f8fe0b858..2850efe58be9 100644 --- a/python/ray/serve/tests/unit/test_replica_quiesce.py +++ b/python/ray/serve/tests/unit/test_replica_quiesce.py @@ -1,6 +1,7 @@ import asyncio import sys -from unittest.mock import MagicMock +import time +from unittest.mock import MagicMock, patch import pytest @@ -50,6 +51,7 @@ def _make_shutdown_fake( grpc_task=None, internal_grpc_port=12345, grace_period_s: float = 1.0, + is_direct_ingress: bool = False, ): """Builds a minimal stand-in for `Replica` for `perform_graceful_shutdown`. @@ -62,6 +64,9 @@ def _make_shutdown_fake( fake._quiescing = False fake._user_callable_initialized = initialized fake._ingress = False + # Set explicitly: on a MagicMock this property would otherwise be truthy and + # silently select the behind-HAProxy drain path. + fake._is_direct_ingress = is_direct_ingress fake._deployment_config.graceful_shutdown_timeout_s = grace_period_s fake._direct_ingress_http_server = ( FakeUvicornServer(events) if with_http_server else None @@ -74,7 +79,7 @@ def _make_shutdown_fake( fake._internal_grpc_port = internal_grpc_port fake._server = FakeGrpcServer(events, "inter_deployment_stop") - async def drain(min_draining_period_s): + async def drain(min_draining_period_s=0.0): # Quiescing must only start AFTER the drain: during the drain the # replica must keep serving normally. assert fake._quiescing is False @@ -82,6 +87,12 @@ async def drain(min_draining_period_s): fake._drain_ongoing_requests = drain + async def drain_behind_haproxy(min_draining_period_s): + assert fake._quiescing is False + events.append(("drain_behind_haproxy", min_draining_period_s)) + + fake._drain_behind_haproxy = drain_behind_haproxy + async def shutdown(): events.append(("shutdown",)) @@ -230,6 +241,89 @@ async def test_uninitialized_replica_skips_drain_and_server_stops(self): assert events == [("shutdown",)] + @pytest.mark.asyncio + async def test_behind_haproxy_uses_two_phase_drain(self): + """Behind HAProxy the two-phase drain is used instead of serving for as + long as requests keep arriving.""" + fake, events = _make_shutdown_fake(is_direct_ingress=True) + + with patch("ray.serve._private.replica.RAY_SERVE_ENABLE_HA_PROXY", True): + await Replica.perform_graceful_shutdown(fake) + + assert [e[0] for e in events] == [ + "drain_behind_haproxy", + "inter_deployment_stop", + "shutdown", + ] + + @pytest.mark.asyncio + async def test_without_haproxy_keeps_serving_until_drained(self): + """Without a retrying party in front, the replica keeps accepting for + the whole drain: refusing would surface to the client.""" + fake, events = _make_shutdown_fake(is_direct_ingress=True) + + with patch("ray.serve._private.replica.RAY_SERVE_ENABLE_HA_PROXY", False): + await Replica.perform_graceful_shutdown(fake) + + assert [e[0] for e in events] == [ + "drain", + "inter_deployment_stop", + "shutdown", + ] + + +class TestDrainBehindHAProxy: + @staticmethod + def _make_fake(): + events = [] + fake = MagicMock() + fake._direct_ingress_http_server = FakeUvicornServer(events) + fake._stop_accepting_direct_ingress = ( + lambda: Replica._stop_accepting_direct_ingress(fake) + ) + + async def drain(min_draining_period_s=0.0): + events.append(("drain", min_draining_period_s)) + + fake._drain_ongoing_requests = drain + return fake, events + + @pytest.mark.asyncio + async def test_closes_listeners_before_waiting_for_in_flight(self): + """The listeners close first, so a request arriving from a stale router + is refused (and retried elsewhere) rather than admitted and then severed + when the replica exits.""" + fake, events = self._make_fake() + + await Replica._drain_behind_haproxy(fake, 0.0) + + assert [e[0] for e in events] == ["http_should_exit", "drain"] + # The minimum period is already spent; it must not be applied again. + assert events[-1] == ("drain", 0.0) + + @pytest.mark.asyncio + async def test_serves_the_full_deregistration_window_first(self): + """The minimum draining period is honoured in full: listeners stay open + until load balancers have had time to deregister this replica.""" + fake, events = self._make_fake() + + start = time.monotonic() + await Replica._drain_behind_haproxy(fake, 0.05) + elapsed = time.monotonic() - start + + assert elapsed >= 0.05 + assert [e[0] for e in events] == ["http_should_exit", "drain"] + + @pytest.mark.asyncio + async def test_no_http_server_is_a_no_op(self): + """A replica without a direct ingress HTTP server still drains.""" + fake, events = self._make_fake() + fake._direct_ingress_http_server = None + + await Replica._drain_behind_haproxy(fake, 0.0) + + assert [e[0] for e in events] == ["drain"] + if __name__ == "__main__": sys.exit(pytest.main(["-v", "-s", __file__])) From c72443a252e97747fca938870f254daa9c50531f Mon Sep 17 00:00:00 2001 From: harshit-anyscale Date: Wed, 26 Aug 2026 16:42:30 +0530 Subject: [PATCH 2/5] Signed-off-by: harshit-anyscale --- python/ray/serve/_private/replica.py | 9 +++------ .../serve/tests/unit/test_replica_quiesce.py | 17 +++++++++++------ 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/python/ray/serve/_private/replica.py b/python/ray/serve/_private/replica.py index f729e3858895..f738f1d584f1 100644 --- a/python/ray/serve/_private/replica.py +++ b/python/ray/serve/_private/replica.py @@ -1226,10 +1226,6 @@ def __init__( def max_ongoing_requests(self) -> int: return self._deployment_config.max_ongoing_requests - @property - def _is_direct_ingress(self) -> bool: - return self._ingress and RAY_SERVE_ENABLE_DIRECT_INGRESS - def get_num_ongoing_requests(self) -> int: return self._metrics_manager.get_num_ongoing_requests() + len( self._reserved_slots @@ -2182,12 +2178,13 @@ def remaining_grace_s() -> float: # In direct ingress mode, hold the replica open at least # RAY_SERVE_DIRECT_INGRESS_MIN_DRAINING_PERIOD_S so load balancers can # deregister it; the drain also waits for in-flight requests. + is_direct_ingress = RAY_SERVE_ENABLE_DIRECT_INGRESS and self._ingress min_draining_period_s = ( RAY_SERVE_DIRECT_INGRESS_MIN_DRAINING_PERIOD_S - if self._is_direct_ingress + if is_direct_ingress else 0.0 ) - if self._is_direct_ingress and RAY_SERVE_ENABLE_HA_PROXY: + if is_direct_ingress and RAY_SERVE_ENABLE_HA_PROXY: # HAProxy retries refused connections, so stop accepting once # the deregistration window is over. await self._drain_behind_haproxy(min_draining_period_s) diff --git a/python/ray/serve/tests/unit/test_replica_quiesce.py b/python/ray/serve/tests/unit/test_replica_quiesce.py index 2850efe58be9..14ab49633770 100644 --- a/python/ray/serve/tests/unit/test_replica_quiesce.py +++ b/python/ray/serve/tests/unit/test_replica_quiesce.py @@ -63,10 +63,11 @@ def _make_shutdown_fake( fake._shutting_down = False fake._quiescing = False fake._user_callable_initialized = initialized - fake._ingress = False - # Set explicitly: on a MagicMock this property would otherwise be truthy and - # silently select the behind-HAProxy drain path. - fake._is_direct_ingress = is_direct_ingress + # Set explicitly: on a MagicMock this would otherwise be truthy and silently + # select the behind-HAProxy drain path. Direct ingress is derived from + # `_ingress` + RAY_SERVE_ENABLE_DIRECT_INGRESS, so the tests that want the + # direct-ingress path patch that flag on as well. + fake._ingress = is_direct_ingress fake._deployment_config.graceful_shutdown_timeout_s = grace_period_s fake._direct_ingress_http_server = ( FakeUvicornServer(events) if with_http_server else None @@ -247,7 +248,9 @@ async def test_behind_haproxy_uses_two_phase_drain(self): long as requests keep arriving.""" fake, events = _make_shutdown_fake(is_direct_ingress=True) - with patch("ray.serve._private.replica.RAY_SERVE_ENABLE_HA_PROXY", True): + with patch( + "ray.serve._private.replica.RAY_SERVE_ENABLE_DIRECT_INGRESS", True + ), patch("ray.serve._private.replica.RAY_SERVE_ENABLE_HA_PROXY", True): await Replica.perform_graceful_shutdown(fake) assert [e[0] for e in events] == [ @@ -262,7 +265,9 @@ async def test_without_haproxy_keeps_serving_until_drained(self): the whole drain: refusing would surface to the client.""" fake, events = _make_shutdown_fake(is_direct_ingress=True) - with patch("ray.serve._private.replica.RAY_SERVE_ENABLE_HA_PROXY", False): + with patch( + "ray.serve._private.replica.RAY_SERVE_ENABLE_DIRECT_INGRESS", True + ), patch("ray.serve._private.replica.RAY_SERVE_ENABLE_HA_PROXY", False): await Replica.perform_graceful_shutdown(fake) assert [e[0] for e in events] == [ From 9635735a2ca805b5f353f686159a11057c5b00d3 Mon Sep 17 00:00:00 2001 From: harshit-anyscale Date: Tue, 1 Sep 2026 12:31:56 +0530 Subject: [PATCH 3/5] [serve] Check ongoing requests before sleeping in the post-window drain Review follow-up. `_drain_ongoing_requests` sleeps `graceful_shutdown_wait_loop_s` before it first counts ongoing requests. That is fine when it also owns the draining period, but in the two-phase drain the window is already spent, so an idle replica waits one more loop for nothing. With `graceful_shutdown_wait_loop_s=10` that overruns the shutdown budget: the deregistration window ends at t=30, the drain then sleeps until t=40, and the controller force-kills at t=35 (the ingress floor of `RAY_SERVE_DIRECT_INGRESS_MIN_DRAINING_PERIOD_S` + `RAY_SERVE_DIRECT_INGRESS_SHUTDOWN_BUFFER_S`), so the quiesce and the graceful server shutdown never run. Add `check_immediately` to count before the first sleep, and pass it from the two-phase drain. The default path is unchanged. Co-Authored-By: Claude Opus 5 Signed-off-by: harshit-anyscale --- python/ray/serve/_private/replica.py | 27 ++++++- .../serve/tests/unit/test_replica_quiesce.py | 75 ++++++++++++++----- 2 files changed, 81 insertions(+), 21 deletions(-) diff --git a/python/ray/serve/_private/replica.py b/python/ray/serve/_private/replica.py index f738f1d584f1..5ebd688f6ff1 100644 --- a/python/ray/serve/_private/replica.py +++ b/python/ray/serve/_private/replica.py @@ -2071,18 +2071,34 @@ def release() -> None: finally: release() - async def _drain_ongoing_requests(self, min_draining_period_s: float = 0.0): + async def _drain_ongoing_requests( + self, + min_draining_period_s: float = 0.0, + check_immediately: bool = False, + ): """Wait until the minimum draining period has elapsed and no ongoing requests remain. The minimum draining period gives load balancers time to deregister this replica; a request admitted during it becomes ongoing and is waited for like any other. + + Args: + min_draining_period_s: keep waiting until at least this long has + passed, even if no requests are ongoing. + check_immediately: count ongoing requests before the first + `graceful_shutdown_wait_loop_s` sleep instead of after it. Use + this when the caller already waited out the draining period, so + an idle replica does not spend another wait loop here. """ wait_loop_period_s = self._deployment_config.graceful_shutdown_wait_loop_s deadline = time.monotonic() + min_draining_period_s + skip_sleep = check_immediately while True: - await asyncio.sleep(wait_loop_period_s) + if skip_sleep: + skip_sleep = False + else: + await asyncio.sleep(wait_loop_period_s) num_ongoing_requests = self.get_num_ongoing_requests() min_period_remaining_s = deadline - time.monotonic() @@ -2140,8 +2156,11 @@ async def _drain_behind_haproxy(self, min_draining_period_s: float) -> None: ) self._stop_accepting_direct_ingress() - # Phase 3: wait for in-flight requests (window already served). - await self._drain_ongoing_requests() + # Phase 3: wait for in-flight requests. Check the count before + # sleeping: the window is already spent, and another + # graceful_shutdown_wait_loop_s here can run past the controller's + # force-kill deadline, which would skip the quiesce below. + await self._drain_ongoing_requests(check_immediately=True) async def shutdown(self): try: diff --git a/python/ray/serve/tests/unit/test_replica_quiesce.py b/python/ray/serve/tests/unit/test_replica_quiesce.py index 14ab49633770..60df34c862e7 100644 --- a/python/ray/serve/tests/unit/test_replica_quiesce.py +++ b/python/ray/serve/tests/unit/test_replica_quiesce.py @@ -80,9 +80,9 @@ def _make_shutdown_fake( fake._internal_grpc_port = internal_grpc_port fake._server = FakeGrpcServer(events, "inter_deployment_stop") - async def drain(min_draining_period_s=0.0): - # Quiescing must only start AFTER the drain: during the drain the - # replica must keep serving normally. + async def drain(min_draining_period_s=0.0, check_immediately=False): + # The drain runs before quiescing, so the replica keeps serving + # normally while it drains. assert fake._quiescing is False events.append(("drain", min_draining_period_s)) @@ -244,8 +244,7 @@ async def test_uninitialized_replica_skips_drain_and_server_stops(self): @pytest.mark.asyncio async def test_behind_haproxy_uses_two_phase_drain(self): - """Behind HAProxy the two-phase drain is used instead of serving for as - long as requests keep arriving.""" + """With HAProxy enabled, shutdown runs the two-phase drain.""" fake, events = _make_shutdown_fake(is_direct_ingress=True) with patch( @@ -261,8 +260,11 @@ async def test_behind_haproxy_uses_two_phase_drain(self): @pytest.mark.asyncio async def test_without_haproxy_keeps_serving_until_drained(self): - """Without a retrying party in front, the replica keeps accepting for - the whole drain: refusing would surface to the client.""" + """With HAProxy disabled, shutdown runs the plain drain. + + Nothing in front would retry a refused connection, so the replica + keeps accepting for the whole drain. + """ fake, events = _make_shutdown_fake(is_direct_ingress=True) with patch( @@ -287,29 +289,34 @@ def _make_fake(): lambda: Replica._stop_accepting_direct_ingress(fake) ) - async def drain(min_draining_period_s=0.0): - events.append(("drain", min_draining_period_s)) + async def drain(min_draining_period_s=0.0, check_immediately=False): + events.append(("drain", min_draining_period_s, check_immediately)) fake._drain_ongoing_requests = drain return fake, events @pytest.mark.asyncio async def test_closes_listeners_before_waiting_for_in_flight(self): - """The listeners close first, so a request arriving from a stale router - is refused (and retried elsewhere) rather than admitted and then severed - when the replica exits.""" + """The HTTP listener closes before we wait for in-flight requests. + + A request that arrives after that is refused, so the caller can retry + it elsewhere instead of having it cut off when the replica exits. + """ fake, events = self._make_fake() await Replica._drain_behind_haproxy(fake, 0.0) assert [e[0] for e in events] == ["http_should_exit", "drain"] - # The minimum period is already spent; it must not be applied again. - assert events[-1] == ("drain", 0.0) + # Phase 1 already waited, so don't wait again, and check the request + # count before sleeping (see `check_immediately`). + assert events[-1] == ("drain", 0.0, True) @pytest.mark.asyncio async def test_serves_the_full_deregistration_window_first(self): - """The minimum draining period is honoured in full: listeners stay open - until load balancers have had time to deregister this replica.""" + """The listener stays open for the whole draining period. + + That is the window load balancers need to deregister the replica. + """ fake, events = self._make_fake() start = time.monotonic() @@ -321,7 +328,7 @@ async def test_serves_the_full_deregistration_window_first(self): @pytest.mark.asyncio async def test_no_http_server_is_a_no_op(self): - """A replica without a direct ingress HTTP server still drains.""" + """A replica with no direct ingress HTTP server still drains.""" fake, events = self._make_fake() fake._direct_ingress_http_server = None @@ -330,5 +337,39 @@ async def test_no_http_server_is_a_no_op(self): assert [e[0] for e in events] == ["drain"] +class TestDrainOngoingRequests: + @staticmethod + def _make_fake(wait_loop_period_s: float, num_ongoing: int = 0): + fake = MagicMock() + fake._deployment_config.graceful_shutdown_wait_loop_s = wait_loop_period_s + fake.get_num_ongoing_requests = lambda: num_ongoing + return fake + + @pytest.mark.asyncio + async def test_check_immediately_skips_the_first_sleep(self): + """With the flag set, an idle replica returns without sleeping first. + + Callers that already waited out the draining period use this: another + wait loop can overrun the controller's force-kill deadline, which + would cut the shutdown short before the servers close gracefully. + """ + fake = self._make_fake(wait_loop_period_s=10) + + start = time.monotonic() + await Replica._drain_ongoing_requests(fake, check_immediately=True) + + assert time.monotonic() - start < 1 + + @pytest.mark.asyncio + async def test_sleeps_before_the_first_check_by_default(self): + """Without the flag, the request count is checked after a wait loop.""" + fake = self._make_fake(wait_loop_period_s=0.05) + + start = time.monotonic() + await Replica._drain_ongoing_requests(fake) + + assert time.monotonic() - start >= 0.05 + + if __name__ == "__main__": sys.exit(pytest.main(["-v", "-s", __file__])) From 945c3df3d5bdca5cd2e362f8ff27f160f3a042c2 Mon Sep 17 00:00:00 2001 From: harshit-anyscale Date: Wed, 2 Sep 2026 16:56:01 +0530 Subject: [PATCH 4/5] [serve] Wait for zero replicas before the second upscale in test_e2e_intermediate_downscaling The test used `check_num_replicas_lte(target=1)` as its cue that the downscale had settled, then cleared the signal and sent 50 more requests. But 1 includes the replica the autoscaler is about to remove, and a replica stays in the router's target list until the controller's stop actually reaches it. That leaves a window (~130ms in CI) where the deployment's target is already 0 but the last replica is still routable. Requests sent in that window are admitted -- `max_ongoing_requests=1000`, so one replica takes all 50 -- and then block forever on the signal the test just cleared. The replica is force-killed at the end of its grace period and the requests die with it. Nothing then reports that demand: replica metrics are filtered to running replicas, and the requests were assigned rather than queued, so the autoscaler reads 0 ongoing requests and never scales back up. The test waits out its 30s timeout. Waiting for the replicas to actually be gone removes the window. This is the same condition the test already uses for its final downscale check, and it is strictly stronger than `<= 1`. Previously this was masked by shutdown latency: idle replicas spent one `graceful_shutdown_wait_loop_s` before their first drain check, so by the time the count dropped the last replica had finished dying. Co-Authored-By: Claude Opus 5 Signed-off-by: harshit-anyscale --- python/ray/serve/tests/test_autoscaling_policy.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/ray/serve/tests/test_autoscaling_policy.py b/python/ray/serve/tests/test_autoscaling_policy.py index 5af9380a6015..7ad60991dbec 100644 --- a/python/ray/serve/tests/test_autoscaling_policy.py +++ b/python/ray/serve/tests/test_autoscaling_policy.py @@ -776,7 +776,9 @@ async def __call__(self): wait_for_condition(check_num_replicas_gte, name="A", target=20, timeout=30) signal.send.remote() - wait_for_condition(check_num_replicas_lte, name="A", target=1, timeout=30) + # Wait for zero, not <= 1: the last replica stays routable until the + # controller's stop reaches it, and would admit the requests below. + wait_for_condition(check_num_replicas_eq, name="A", target=0, timeout=30) signal.send.remote(clear=True) [handle.remote() for _ in range(50)] From 43eb10153b4d6db9a485219239989c1786861a55 Mon Sep 17 00:00:00 2001 From: harshit-anyscale Date: Thu, 3 Sep 2026 12:58:56 +0530 Subject: [PATCH 5/5] [serve] Make perform_graceful_shutdown idempotent The controller re-issues the stop when it restarts, so a replica can get a second `perform_graceful_shutdown` while the first is still running. Nothing prevented re-entry, and `call_destructor` deliberately runs the user's `__del__` only once -- so the second pass skipped the destructor, returned promptly, and the controller recorded a clean shutdown while the first pass was still blocked inside `__del__`. The replica is then reported stopped and the application deletes with user cleanup still pending, which is exactly what test_recover_deleting_application checks. Run the shutdown once and have later calls await that same run, shielded so a cancelled caller can't abort it. The controller now keeps seeing the replica as stopping until the grace period expires and force-kills it. Previously the second pass spent one `graceful_shutdown_wait_loop_s` in the drain before reaching the destructor, which was long enough to hide this. Co-Authored-By: Claude Opus 5 Signed-off-by: harshit-anyscale --- python/ray/serve/_private/replica.py | 21 +++++ .../serve/tests/unit/test_replica_quiesce.py | 76 +++++++++++++++++-- 2 files changed, 90 insertions(+), 7 deletions(-) diff --git a/python/ray/serve/_private/replica.py b/python/ray/serve/_private/replica.py index 21affc96d930..9c86da5eec82 100644 --- a/python/ray/serve/_private/replica.py +++ b/python/ray/serve/_private/replica.py @@ -1237,6 +1237,10 @@ def __init__( # requests are then rejected (the router retries them elsewhere). self._quiescing = False + # Tracks the in-progress graceful shutdown so repeated calls await the + # same one (see `perform_graceful_shutdown`). + self._graceful_shutdown_task: Optional[asyncio.Task] = None + self._num_queued_requests = 0 self._reserved_slots: Set[str] = set() @@ -2198,6 +2202,23 @@ async def shutdown(self): await self._metrics_manager.shutdown() async def perform_graceful_shutdown(self): + """Shut down gracefully, at most once. + + The controller re-issues the stop when it restarts, so this can be + called more than once for the same replica. Later calls must await the + first rather than run a second pass: the user's destructor only runs + once, so a second pass would skip it and report a clean shutdown while + the first is still inside `__del__`. + """ + if self._graceful_shutdown_task is None: + self._graceful_shutdown_task = self._event_loop.create_task( + self._perform_graceful_shutdown() + ) + + # Shielded so a cancelled caller doesn't abort the shutdown itself. + await asyncio.shield(self._graceful_shutdown_task) + + async def _perform_graceful_shutdown(self): self._shutting_down = True # Shutdown budget, mirroring the controller's force-kill deadline (see diff --git a/python/ray/serve/tests/unit/test_replica_quiesce.py b/python/ray/serve/tests/unit/test_replica_quiesce.py index 60df34c862e7..65a2c9443fc2 100644 --- a/python/ray/serve/tests/unit/test_replica_quiesce.py +++ b/python/ray/serve/tests/unit/test_replica_quiesce.py @@ -53,7 +53,7 @@ def _make_shutdown_fake( grace_period_s: float = 1.0, is_direct_ingress: bool = False, ): - """Builds a minimal stand-in for `Replica` for `perform_graceful_shutdown`. + """Builds a minimal stand-in for `Replica` for `_perform_graceful_shutdown`. Returns the fake and the ordered list of events it records. """ @@ -170,7 +170,7 @@ async def test_quiesces_and_stops_servers_in_order(self): grace_period_s=1.0, ) - await Replica.perform_graceful_shutdown(fake) + await Replica._perform_graceful_shutdown(fake) assert fake._shutting_down is True assert fake._quiescing is True @@ -205,7 +205,7 @@ async def test_http_graceful_close_timeout_falls_back_to_cancel(self): grace_period_s=0.05, ) - await Replica.perform_graceful_shutdown(fake) + await Replica._perform_graceful_shutdown(fake) # `asyncio.wait_for` cancels the awaited task on timeout. assert http_task.cancelled() @@ -220,7 +220,7 @@ async def test_abrupt_cancel_when_server_object_missing(self): http_task = MagicMock() fake, events = _make_shutdown_fake(http_task=http_task) - await Replica.perform_graceful_shutdown(fake) + await Replica._perform_graceful_shutdown(fake) assert http_task.cancel.called assert [e[0] for e in events] == [ @@ -238,7 +238,7 @@ async def test_uninitialized_replica_skips_drain_and_server_stops(self): internal_grpc_port=None, ) - await Replica.perform_graceful_shutdown(fake) + await Replica._perform_graceful_shutdown(fake) assert events == [("shutdown",)] @@ -250,7 +250,7 @@ async def test_behind_haproxy_uses_two_phase_drain(self): with patch( "ray.serve._private.replica.RAY_SERVE_ENABLE_DIRECT_INGRESS", True ), patch("ray.serve._private.replica.RAY_SERVE_ENABLE_HA_PROXY", True): - await Replica.perform_graceful_shutdown(fake) + await Replica._perform_graceful_shutdown(fake) assert [e[0] for e in events] == [ "drain_behind_haproxy", @@ -270,7 +270,7 @@ async def test_without_haproxy_keeps_serving_until_drained(self): with patch( "ray.serve._private.replica.RAY_SERVE_ENABLE_DIRECT_INGRESS", True ), patch("ray.serve._private.replica.RAY_SERVE_ENABLE_HA_PROXY", False): - await Replica.perform_graceful_shutdown(fake) + await Replica._perform_graceful_shutdown(fake) assert [e[0] for e in events] == [ "drain", @@ -371,5 +371,67 @@ async def test_sleeps_before_the_first_check_by_default(self): assert time.monotonic() - start >= 0.05 +class TestGracefulShutdownIsIdempotent: + """The controller re-issues the stop when it restarts, so the shutdown + must run once and later calls must await that same run.""" + + @staticmethod + def _make_fake(shutdown_body): + fake = MagicMock() + fake._graceful_shutdown_task = None + fake._event_loop = asyncio.get_event_loop() + fake._perform_graceful_shutdown = shutdown_body + return fake + + @pytest.mark.asyncio + async def test_second_call_does_not_start_a_second_shutdown(self): + calls = [] + + async def body(): + calls.append("shutdown") + await asyncio.sleep(0.05) + + fake = self._make_fake(body) + await asyncio.gather( + Replica.perform_graceful_shutdown(fake), + Replica.perform_graceful_shutdown(fake), + ) + + assert calls == ["shutdown"] + + @pytest.mark.asyncio + async def test_second_call_waits_for_the_first_to_finish(self): + finished = [] + + async def body(): + await asyncio.sleep(0.05) + finished.append(True) + + fake = self._make_fake(body) + first = asyncio.ensure_future(Replica.perform_graceful_shutdown(fake)) + await asyncio.sleep(0) # let the first call create the task + + # A caller arriving mid-shutdown must not return before it completes. + await Replica.perform_graceful_shutdown(fake) + assert finished == [True] + await first + + @pytest.mark.asyncio + async def test_cancelled_caller_does_not_abort_the_shutdown(self): + finished = [] + + async def body(): + await asyncio.sleep(0.05) + finished.append(True) + + fake = self._make_fake(body) + caller = asyncio.ensure_future(Replica.perform_graceful_shutdown(fake)) + await asyncio.sleep(0) + caller.cancel() + + await asyncio.sleep(0.1) + assert finished == [True] + + if __name__ == "__main__": sys.exit(pytest.main(["-v", "-s", __file__]))