chore: import upstream snapshot with attribution

python/ray/serve/tests/test_consistent_hash_router.py

@@ -0,0 +1,354 @@
+import sys
+
+import grpc
+import httpx
+import pytest
+
+import ray
+from ray import serve
+from ray._common.test_utils import SignalActor, wait_for_condition
+from ray.serve._private.constants import SERVE_NAMESPACE
+from ray.serve._private.test_utils import check_running, get_application_url
+from ray.serve.config import RequestRouterConfig, gRPCOptions
+from ray.serve.context import _get_internal_replica_context
+from ray.serve.generated import serve_pb2, serve_pb2_grpc
+
+ROUTER_CLASS = "ray.serve.experimental.consistent_hash_router:ConsistentHashRouter"
+
+
+class TestConsistentHashRouting:
+    def test_same_session_sticky_to_same_replica(self, serve_instance):
+        @serve.deployment(
+            request_router_config=RequestRouterConfig(
+                request_router_class=ROUTER_CLASS,
+                request_router_kwargs={
+                    "num_virtual_nodes": 100,
+                    "num_fallback_replicas": 2,
+                },
+                initial_backoff_s=0.01,
+                backoff_multiplier=2.0,
+                max_backoff_s=0.1,
+            ),
+            num_replicas=4,
+            max_ongoing_requests=5,
+            ray_actor_options={"num_cpus": 0},
+        )
+        class App:
+            def __init__(self):
+                self.unique_id = _get_internal_replica_context().replica_id.unique_id
+
+            async def __call__(self):
+                return self.unique_id
+
+        handle = serve.run(App.bind())
+        wait_for_condition(check_running, timeout=30)
+
+        session = "session-123"
+        replicas = [
+            handle.options(session_id=session).remote().result(timeout_s=10)
+            for _ in range(20)
+        ]
+        assert len(set(replicas)) == 1
+
+    def test_different_sessions_spread_across_replicas(self, serve_instance):
+        @serve.deployment(
+            request_router_config=RequestRouterConfig(
+                request_router_class=ROUTER_CLASS,
+                request_router_kwargs={
+                    "num_virtual_nodes": 100,
+                    "num_fallback_replicas": 2,
+                },
+                initial_backoff_s=0.01,
+                backoff_multiplier=2.0,
+                max_backoff_s=0.1,
+            ),
+            num_replicas=4,
+            max_ongoing_requests=5,
+            ray_actor_options={"num_cpus": 0},
+        )
+        class App:
+            def __init__(self):
+                self.unique_id = _get_internal_replica_context().replica_id.unique_id
+
+            async def __call__(self):
+                return self.unique_id
+
+        handle = serve.run(App.bind())
+        wait_for_condition(check_running, timeout=30)
+
+        replicas = [
+            handle.options(session_id=f"session_{i}").remote().result(timeout_s=10)
+            for i in range(40)
+        ]
+        # With MurmurHash3, the probability of all 40 sessions landing on
+        # the same replica is negligible.
+        assert len(set(replicas)) > 1
+
+    def test_request_without_session_id_still_routes(self, serve_instance):
+        @serve.deployment(
+            request_router_config=RequestRouterConfig(
+                request_router_class=ROUTER_CLASS,
+                request_router_kwargs={
+                    "num_virtual_nodes": 100,
+                    "num_fallback_replicas": 2,
+                },
+                initial_backoff_s=0.01,
+                backoff_multiplier=2.0,
+                max_backoff_s=0.1,
+            ),
+            num_replicas=3,
+            max_ongoing_requests=5,
+            ray_actor_options={"num_cpus": 0},
+        )
+        class App:
+            def __init__(self):
+                self.unique_id = _get_internal_replica_context().replica_id.unique_id
+
+            async def __call__(self):
+                return self.unique_id
+
+        handle = serve.run(App.bind())
+        wait_for_condition(check_running, timeout=30)
+
+        for _ in range(10):
+            result = handle.remote().result(timeout_s=10)
+            assert isinstance(result, str)
+            assert len(result) > 0
+
+    def test_mixed_sessioned_and_sessionless_traffic(self, serve_instance):
+        @serve.deployment(
+            request_router_config=RequestRouterConfig(
+                request_router_class=ROUTER_CLASS,
+                request_router_kwargs={
+                    "num_virtual_nodes": 100,
+                    "num_fallback_replicas": 2,
+                },
+                initial_backoff_s=0.01,
+                backoff_multiplier=2.0,
+                max_backoff_s=0.1,
+            ),
+            num_replicas=4,
+            max_ongoing_requests=5,
+            ray_actor_options={"num_cpus": 0},
+        )
+        class App:
+            def __init__(self):
+                self.unique_id = _get_internal_replica_context().replica_id.unique_id
+
+            async def __call__(self):
+                return self.unique_id
+
+        handle = serve.run(App.bind())
+        wait_for_condition(check_running, timeout=30)
+
+        sticky_session = "sticky-session"
+        sticky_landings = []
+        for _ in range(10):
+            sticky_landings.append(
+                handle.options(session_id=sticky_session).remote().result(timeout_s=10)
+            )
+            # Interleave a session-less request.
+            handle.remote().result(timeout_s=10)
+
+        assert len(set(sticky_landings)) == 1
+
+
+class TestOverflowToFallback:
+    """
+    When the primary is at max_ongoing_requests, session traffic must
+    overflow to the fallback replica rather than hang.
+    """
+
+    def test_overflow_when_primary_saturated(self, serve_instance):
+        signal_actor_name = "consistent-hash-signal"
+        signal = SignalActor.options(name=signal_actor_name).remote()
+
+        # max_ongoing_requests=1 so a single in-flight request fully
+        # saturates the primary.
+        @serve.deployment(
+            request_router_config=RequestRouterConfig(
+                request_router_class=ROUTER_CLASS,
+                request_router_kwargs={
+                    "num_virtual_nodes": 100,
+                    "num_fallback_replicas": 2,
+                },
+                initial_backoff_s=0.01,
+                backoff_multiplier=2.0,
+                max_backoff_s=0.1,
+            ),
+            num_replicas=3,
+            max_ongoing_requests=1,
+            ray_actor_options={"num_cpus": 0},
+        )
+        class BlockingApp:
+            def __init__(self):
+                self.unique_id = _get_internal_replica_context().replica_id.unique_id
+
+            async def __call__(self):
+                await ray.get_actor(signal_actor_name).wait.remote()
+                return self.unique_id
+
+        handle = serve.run(BlockingApp.bind())
+        wait_for_condition(check_running, timeout=30)
+
+        session = "overflow-session"
+
+        # Fire the first request on the session. It takes the primary and blocks on the signal.
+        first_ref = handle.options(session_id=session).remote()
+        wait_for_condition(
+            lambda: ray.get(signal.cur_num_waiters.remote()) == 1,
+            timeout=10,
+        )
+
+        # Second request on the same session -- primary is at max_ongoing_requests,
+        # so the retry loop should walk to fallback_1 and that replica picks it up.
+        second_ref = handle.options(session_id=session).remote()
+        wait_for_condition(
+            lambda: ray.get(signal.cur_num_waiters.remote()) == 2,
+            timeout=10,
+        )
+
+        # Release both and read the replica ids.
+        ray.get(signal.send.remote())
+        first_replica = first_ref.result(timeout_s=10)
+        second_replica = second_ref.result(timeout_s=10)
+
+        assert first_replica != second_replica
+
+    def test_sticky_returns_to_primary_after_drain(self, serve_instance):
+        signal_actor_name = "consistent-hash-signal-drain"
+        signal = SignalActor.options(name=signal_actor_name).remote()
+
+        @serve.deployment(
+            request_router_config=RequestRouterConfig(
+                request_router_class=ROUTER_CLASS,
+                request_router_kwargs={
+                    "num_virtual_nodes": 100,
+                    "num_fallback_replicas": 2,
+                },
+                initial_backoff_s=0.01,
+                backoff_multiplier=2.0,
+                max_backoff_s=0.1,
+            ),
+            num_replicas=3,
+            max_ongoing_requests=1,
+            ray_actor_options={"num_cpus": 0},
+        )
+        class BlockingApp:
+            def __init__(self):
+                self.unique_id = _get_internal_replica_context().replica_id.unique_id
+
+            async def __call__(self):
+                await ray.get_actor(signal_actor_name).wait.remote()
+                return self.unique_id
+
+        handle = serve.run(BlockingApp.bind())
+        wait_for_condition(check_running, timeout=30)
+
+        session = "returning-session"
+
+        # Fire one blocking request to identify the primary.
+        first_ref = handle.options(session_id=session).remote()
+        wait_for_condition(
+            lambda: ray.get(signal.cur_num_waiters.remote()) == 1,
+            timeout=10,
+        )
+        ray.get(signal.send.remote())
+        primary = first_ref.result(timeout_s=10)
+
+        # Now fire several non-blocking follow-up requests on the same session.
+        # They should all land on the primary because the ring is unchanged and
+        # the primary is free.
+        follow_ups = [
+            handle.options(session_id=session).remote().result(timeout_s=10)
+            for _ in range(5)
+        ]
+        assert all(r == primary for r in follow_ups)
+
+
+class TestProtocolStickiness:
+    def test_http_same_session_sticky(self, serve_instance):
+        @serve.deployment(
+            request_router_config=RequestRouterConfig(
+                request_router_class=ROUTER_CLASS,
+                request_router_kwargs={
+                    "num_virtual_nodes": 100,
+                    "num_fallback_replicas": 2,
+                },
+            ),
+            num_replicas=4,
+            ray_actor_options={"num_cpus": 0},
+        )
+        class App:
+            async def __call__(self, request):
+                return _get_internal_replica_context().replica_id.unique_id
+
+        serve.run(App.bind())
+        wait_for_condition(check_running, timeout=30)
+
+        url = get_application_url()
+        session_replicas = {
+            httpx.get(url, headers={"x-session-id": "sess_http_42"}, timeout=10).text
+            for _ in range(15)
+        }
+        other_replicas = {
+            httpx.get(url, headers={"x-session-id": "sess_http_99"}, timeout=10).text
+            for _ in range(15)
+        }
+        assert len(session_replicas) == 1, f"sess_http_42 drifted: {session_replicas}"
+        assert len(other_replicas) == 1, f"sess_http_99 drifted: {other_replicas}"
+
+    def test_grpc_same_session_sticky(self, ray_cluster):
+        cluster = ray_cluster
+        cluster.add_node(num_cpus=2)
+        cluster.connect(namespace=SERVE_NAMESPACE)
+
+        serve.start(
+            grpc_options=gRPCOptions(
+                port=9000,
+                grpc_servicer_functions=[
+                    "ray.serve.generated.serve_pb2_grpc."
+                    "add_UserDefinedServiceServicer_to_server",
+                ],
+            ),
+        )
+
+        @serve.deployment(
+            request_router_config=RequestRouterConfig(
+                request_router_class=ROUTER_CLASS,
+                request_router_kwargs={
+                    "num_virtual_nodes": 100,
+                    "num_fallback_replicas": 2,
+                },
+            ),
+            num_replicas=4,
+            ray_actor_options={"num_cpus": 0},
+        )
+        class GrpcApp:
+            def __call__(self, user_message):
+                return serve_pb2.UserDefinedResponse(
+                    greeting=_get_internal_replica_context().replica_id.unique_id,
+                )
+
+        serve.run(GrpcApp.bind())
+        wait_for_condition(check_running, timeout=30)
+
+        channel = grpc.insecure_channel(get_application_url("gRPC", use_localhost=True))
+        stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
+        req = serve_pb2.UserDefinedMessage(name="x", num=1, foo="y")
+
+        def run(session_id: str) -> set:
+            metadata = (("session_id", session_id),)
+            landings = set()
+            for _ in range(15):
+                resp, call = stub.__call__.with_call(request=req, metadata=metadata)
+                assert call.code() == grpc.StatusCode.OK
+                landings.add(resp.greeting)
+            return landings
+
+        assert len(run("sess_grpc_42")) == 1
+        assert len(run("sess_grpc_99")) == 1
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))