[12/N][HAProxy stability] - Close direct-ingress listeners once the deregistration window ends

python/ray/serve/_private/replica.py

@@ -2100,6 +2100,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()
@@ -2135,12 +2178,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

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`.
 
@@ -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):
+    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
         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",))
 
@@ -230,6 +242,93 @@ 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_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):
+        """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_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):
+            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__]))