diff --git a/python/ray/serve/_private/replica.py b/python/ray/serve/_private/replica.py index 15728e5b7bb9..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() @@ -2089,18 +2093,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() @@ -2118,6 +2138,52 @@ 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. 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: self._user_callable_wrapper.stop_user_loop_watchdog() @@ -2136,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 @@ -2153,12 +2236,19 @@ 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 RAY_SERVE_ENABLE_DIRECT_INGRESS and self._ingress + if is_direct_ingress else 0.0 ) - await self._drain_ongoing_requests(min_draining_period_s) + 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) + 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/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)] diff --git a/python/ray/serve/tests/unit/test_replica_quiesce.py b/python/ray/serve/tests/unit/test_replica_quiesce.py index 331f8fe0b858..65a2c9443fc2 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,8 +51,9 @@ 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`. + """Builds a minimal stand-in for `Replica` for `_perform_graceful_shutdown`. Returns the fake and the ordered list of events it records. """ @@ -61,7 +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 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 @@ -74,14 +80,20 @@ 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): - # 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)) 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",)) @@ -158,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 @@ -193,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() @@ -208,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] == [ @@ -226,10 +238,200 @@ 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",)] + @pytest.mark.asyncio + async def test_behind_haproxy_uses_two_phase_drain(self): + """With HAProxy enabled, shutdown runs the two-phase drain.""" + fake, events = _make_shutdown_fake(is_direct_ingress=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] == [ + "drain_behind_haproxy", + "inter_deployment_stop", + "shutdown", + ] + + @pytest.mark.asyncio + async def test_without_haproxy_keeps_serving_until_drained(self): + """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( + "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] == [ + "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, 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 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"] + # 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 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() + 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 with no 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"] + + +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 + + +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__]))