import asyncio import sys from unittest.mock import MagicMock import pytest from ray.serve._private.common import RequestMetadata from ray.serve._private.replica import Replica def _make_metadata(*, is_direct_ingress: bool) -> RequestMetadata: return RequestMetadata( request_id="test-request", internal_request_id="test-internal-request", is_direct_ingress=is_direct_ingress, ) class FakeUvicornServer: """Stand-in for `uvicorn.Server` that records when `should_exit` is set.""" def __init__(self, events): # Write through __dict__ to avoid recursing into __setattr__. self.__dict__["_events"] = events self.__dict__["should_exit"] = False def __setattr__(self, name, value): if name == "should_exit": self._events.append(("http_should_exit", value)) self.__dict__[name] = value class FakeGrpcServer: """Stand-in for a `grpc.aio` server that records graceful stops.""" def __init__(self, events, name): self._events = events self._name = name async def stop(self, grace): self._events.append((self._name, grace)) def _make_shutdown_fake( *, initialized: bool = True, with_http_server: bool = False, http_task=None, with_grpc_server: bool = False, grpc_task=None, internal_grpc_port=12345, grace_period_s: float = 1.0, ): """Builds a minimal stand-in for `Replica` for `perform_graceful_shutdown`. Returns the fake and the ordered list of events it records. """ events = [] fake = MagicMock() fake._shutting_down = False fake._quiescing = False fake._user_callable_initialized = initialized fake._ingress = False fake._deployment_config.graceful_shutdown_timeout_s = grace_period_s fake._direct_ingress_http_server = ( FakeUvicornServer(events) if with_http_server else None ) fake._direct_ingress_http_server_task = http_task fake._direct_ingress_grpc_server = ( FakeGrpcServer(events, "direct_ingress_grpc_stop") if with_grpc_server else None ) fake._direct_ingress_grpc_server_task = grpc_task 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. assert fake._quiescing is False events.append(("drain", min_draining_period_s)) fake._drain_ongoing_requests = drain async def shutdown(): events.append(("shutdown",)) fake.shutdown = shutdown return fake, events class TestCanAcceptRequestWhileQuiescing: @staticmethod def _bind_quiescing(fake): # _can_accept_request delegates the quiescing decision to # _is_replica_quiescing; wire the real helper onto the fake so it runs # instead of returning a truthy MagicMock. fake._is_replica_quiescing = lambda rm: Replica._is_replica_quiescing(fake, rm) def test_handle_path_rejected_when_quiescing(self): fake = MagicMock() fake._quiescing = True self._bind_quiescing(fake) assert ( Replica._can_accept_request(fake, _make_metadata(is_direct_ingress=False)) is False ) def test_direct_ingress_not_rejected_when_quiescing(self): """Direct ingress requests must NOT be rejected while quiescing. The HTTP/gRPC servers are shut down gracefully during quiescing, so any direct ingress request that still arrives is served to completion; rejecting it would surface a client-visible error instead of a safe retry. """ fake = MagicMock() fake._quiescing = True fake.max_queued_requests = -1 fake._num_queued_requests = 0 self._bind_quiescing(fake) assert ( Replica._can_accept_request(fake, _make_metadata(is_direct_ingress=True)) is True ) def test_handle_path_accepted_when_not_quiescing(self): fake = MagicMock() fake._quiescing = False fake._semaphore.locked.return_value = False self._bind_quiescing(fake) assert ( Replica._can_accept_request(fake, _make_metadata(is_direct_ingress=False)) is True ) class TestPerformGracefulShutdown: @pytest.mark.asyncio async def test_quiesces_and_stops_servers_in_order(self): """Drain → quiesce → graceful server stops → shutdown, in order.""" loop = asyncio.get_running_loop() http_task = loop.create_future() # Simulate the uvicorn serve task exiting promptly once # `should_exit` is set. http_task.set_result(None) grpc_task = MagicMock() fake, events = _make_shutdown_fake( with_http_server=True, http_task=http_task, with_grpc_server=True, grpc_task=grpc_task, grace_period_s=1.0, ) await Replica.perform_graceful_shutdown(fake) assert fake._shutting_down is True assert fake._quiescing is True assert [e[0] for e in events] == [ "drain", "http_should_exit", "direct_ingress_grpc_stop", "inter_deployment_stop", "shutdown", ] # The grace passed to each stop is the REMAINING shutdown budget at # that step (deadline-based), so it must be positive and within the # configured budget. for name, *args in events: if name.endswith("_stop"): assert 0.0 < args[0] <= 1.0 # The direct ingress gRPC server task is cancelled after the # graceful stop completes. assert grpc_task.cancel.called @pytest.mark.asyncio async def test_http_graceful_close_timeout_falls_back_to_cancel(self): """If the HTTP server doesn't exit within the grace period, the server task is cancelled (the previous abrupt behavior).""" loop = asyncio.get_running_loop() # A serve task that never exits on its own. http_task = loop.create_future() fake, events = _make_shutdown_fake( with_http_server=True, http_task=http_task, grace_period_s=0.05, ) await Replica.perform_graceful_shutdown(fake) # `asyncio.wait_for` cancels the awaited task on timeout. assert http_task.cancelled() # Shutdown still completes despite the timeout. assert fake._quiescing is True assert events[-1] == ("shutdown",) @pytest.mark.asyncio async def test_abrupt_cancel_when_server_object_missing(self): """If only the server task exists (no server object), fall back to the abrupt cancel.""" http_task = MagicMock() fake, events = _make_shutdown_fake(http_task=http_task) await Replica.perform_graceful_shutdown(fake) assert http_task.cancel.called assert [e[0] for e in events] == [ "drain", "inter_deployment_stop", "shutdown", ] @pytest.mark.asyncio async def test_uninitialized_replica_skips_drain_and_server_stops(self): """A replica that never initialized has no servers to stop and never served traffic, so only the final shutdown step runs.""" fake, events = _make_shutdown_fake( initialized=False, internal_grpc_port=None, ) await Replica.perform_graceful_shutdown(fake) assert events == [("shutdown",)] if __name__ == "__main__": sys.exit(pytest.main(["-v", "-s", __file__]))