[Serve] Rewrite Router to be Embeddable (#12019)

python/ray/serve/api.py

@@ -1,6 +1,5 @@
 import atexit
 from functools import wraps
-import random
 import os
 
 import ray
@@ -9,8 +8,7 @@ from ray.serve.constants import (DEFAULT_HTTP_HOST, DEFAULT_HTTP_PORT,
 from ray.serve.controller import ServeController
 from ray.serve.handle import RayServeHandle
 from ray.serve.utils import (block_until_http_ready, format_actor_name,
-                             get_random_letters, logger, get_node_id_for_actor,
-                             get_conda_env_dir)
+                             get_random_letters, logger, get_conda_env_dir)
 from ray.serve.exceptions import RayServeException
 from ray.serve.config import BackendConfig, ReplicaConfig, BackendMetadata
 from ray.serve.env import CondaEnv
@@ -45,6 +43,13 @@ class Client:
         self._detached = detached
         self._shutdown = False
 
+        # NOTE(simon): Used to cache client.get_handle(endpoint) call. It will
+        # mostly grow in size, it will only shrink when user calls the
+        # .remove_endpoint method. This is fine because we expect the number of
+        # endpoints to be fairly small. However, in case this dictionary does
+        # grow very big, we can replace it with a LRU cache instead.
+        self._handle_cache: Dict[str, ActorHandle] = dict()
+
         # NOTE(edoakes): Need this because the shutdown order isn't guaranteed
         # when the interpreter is exiting so we can't rely on __del__ (it
         # throws a nasty stacktrace).
@@ -117,7 +122,7 @@ class Client:
         if endpoint_name in endpoints:
             methods_old = endpoints[endpoint_name]["methods"]
             route_old = endpoints[endpoint_name]["route"]
-            if methods_old.sort() == methods.sort() and route_old == route:
+            if sorted(methods_old) == sorted(methods) and route_old == route:
                 raise ValueError(
                     "Route '{}' is already registered to endpoint '{}' "
                     "with methods '{}'.  To set the backend for this "
@@ -142,6 +147,8 @@ class Client:
 
         Does not delete any associated backends.
         """
+        if endpoint in self._handle_cache:
+            del self._handle_cache[endpoint]
         ray.get(self._controller.delete_endpoint.remote(endpoint))
 
     @_ensure_connected
@@ -366,23 +373,10 @@ class Client:
                 self._controller.get_all_endpoints.remote()):
             raise KeyError(f"Endpoint '{endpoint_name}' does not exist.")
 
-        routers = list(ray.get(self._controller.get_routers.remote()).values())
-        current_node_id = ray.get_runtime_context().node_id.hex()
-
-        try:
-            router_chosen = next(
-                filter(lambda r: get_node_id_for_actor(r) == current_node_id,
-                       routers))
-        except StopIteration:
-            logger.warning(
-                f"When getting a handle for {endpoint_name}, Serve can't find "
-                "a router on the same node. Serve will use a random router.")
-            router_chosen = random.choice(routers)
-
-        return RayServeHandle(
-            router_chosen,
-            endpoint_name,
-        )
+        if endpoint_name not in self._handle_cache:
+            handle = RayServeHandle(self._controller, endpoint_name, sync=True)
+            self._handle_cache[endpoint_name] = handle
+        return self._handle_cache[endpoint_name]
 
 
 def start(detached: bool = False,

python/ray/serve/controller.py

@@ -45,8 +45,8 @@ NodeId = str
 
 class TrafficPolicy:
     def __init__(self, traffic_dict: Dict[str, float]) -> None:
-        self.traffic_dict = dict()
-        self.shadow_dict = dict()
+        self.traffic_dict: Dict[str, float] = dict()
+        self.shadow_dict: Dict[str, float] = dict()
         self.set_traffic_dict(traffic_dict)
 
     def set_traffic_dict(self, traffic_dict: Dict[str, float]) -> None:
@@ -71,6 +71,9 @@ class TrafficPolicy:
         else:
             self.shadow_dict[backend] = proportion
 
+    def __repr__(self) -> str:
+        return f"<Traffic {self.traffic_dict}; Shadow {self.shadow_dict}>"
+
 
 class BackendInfo(BaseModel):
     # TODO(architkulkarni): Add type hint for worker_class after upgrading
@@ -318,7 +321,6 @@ class ActorStateReconciler:
                         node_resource: 0.01
                     },
                 ).remote(
-                    node_id,
                     http_host,
                     http_port,
                     controller_name=self.controller_name,
@@ -484,7 +486,11 @@ class ServeController:
 
     def notify_replica_handles_changed(self):
         self.long_poll_host.notify_changed(
-            "worker_handles", self.actor_reconciler.backend_replicas)
+            "worker_handles", {
+                backend_tag: list(replica_dict.values())
+                for backend_tag, replica_dict in
+                self.actor_reconciler.backend_replicas.items()
+            })
 
     def notify_traffic_policies_changed(self):
         self.long_poll_host.notify_changed(

python/ray/serve/endpoint_policy.py

@@ -1,15 +1,32 @@
 from abc import ABCMeta, abstractmethod
-import copy
+import random
 from hashlib import sha256
+from functools import lru_cache
+from typing import List
 
 import numpy as np
 
+import ray
 from ray.serve.utils import logger
 
 
+@lru_cache(maxsize=128)
+def deterministic_hash(key: bytes) -> float:
+    """Given an arbitrary bytes, return a deterministic value between 0 and 1.
+
+    Note:
+        This function uses stdlib random number generator because it's faster
+        than numpy's. On a cache miss, the runtime of this function is about
+        ~10us.
+    """
+    bytes_hash = sha256(key).digest()  # should return 32 bytes value
+    int_seed = int.from_bytes(bytes_hash, "little", signed=False)
+    random_state = random.Random(int_seed)
+    return random_state.random()
+
+
 class EndpointPolicy:
     """Defines the interface for a routing policy for a single endpoint.
-
     To add a new routing policy, a class should be defined that provides this
     interface. The class may be stateful, in which case it may also want to
     provide a non-default constructor. However, this state will be lost when
@@ -18,37 +35,29 @@ class EndpointPolicy:
     __metaclass__ = ABCMeta
 
     @abstractmethod
-    def flush(self, endpoint_queue, backend_queues):
-        """Flush the endpoint queue into the given backend queues.
-
-        This method should assign each query in the endpoint_queue to a
-        backend in the backend_queues. Queries are assigned by popping them
-        from the endpoint queue and pushing them onto a backend queue. The
-        method must also return a set of all backend tags so that the caller
-        knows which backend_queues to flush.
+    def assign(self, query) -> List[str]:
+        """Assign a query to a list of backends.
 
         Arguments:
-            endpoint_queue: deque containing queries to assign.
-            backend_queues: Dict(str, deque) mapping backend tags to
-            their corresponding query queues.
-
+            query (ray.serve.router.Query): the incoming query object.
         Returns:
-            Set of backend tags that had queries added to their queues.
+            A list with length >= 1. It should contains the list of backend
+            tags that had queries added to their queues. Ordered by importance.
+            The first value should be the backend to assign and rest values
+            correspond to shadow backends.
         """
-        assigned_backends = set()
-        return assigned_backends
+        raise NotImplementedError()
 
 
 class RandomEndpointPolicy(EndpointPolicy):
     """
     A stateless policy that makes a weighted random decision to map each query
-    to a backend using the specified weights.
-
-    If a shard key is provided in a query, the weighted random selection will
-    be made deterministically based on the hash of the shard key.
+    to a backend using the specified weights. If a shard key is provided in a
+    query, the weighted random selection will be made deterministically based
+    on the hash of the shard key.
     """
 
-    def __init__(self, traffic_policy):
+    def __init__(self, traffic_policy: "ray.serve.controller.TrafficPolicy"):
         self.backends = sorted(traffic_policy.traffic_dict.items())
         self.shadow_backends = list(traffic_policy.shadow_dict.items())
 
@@ -69,33 +78,16 @@ class RandomEndpointPolicy(EndpointPolicy):
 
         return chosen_backend, shadow_backends
 
-    def flush(self, endpoint_queue, backend_queues):
+    def assign(self, query):
         if len(self.backends) == 0:
-            logger.info("No backends to assign traffic to.")
-            return set()
+            raise ValueError("No backends to assign traffic to.")
 
-        assigned_backends = set()
-        while len(endpoint_queue) > 0:
-            query = endpoint_queue.pop()
-            if query.metadata.shard_key is None:
-                rstate = np.random
-            else:
-                sha256_seed = sha256(query.metadata.shard_key.encode("utf-8"))
-                seed = np.frombuffer(sha256_seed.digest(), dtype=np.uint32)
-                # Note(simon): This constructor takes 100+us, maybe cache this?
-                rstate = np.random.RandomState(seed)
+        if query.metadata.shard_key is None:
+            value = np.random.random()
+        else:
+            value = deterministic_hash(
+                query.metadata.shard_key.encode("utf-8"))
 
-            chosen_backend, shadow_backends = self._select_backends(
-                rstate.random())
-
-            assigned_backends.add(chosen_backend)
-            backend_queues[chosen_backend].appendleft(query)
-            if len(shadow_backends) > 0:
-                shadow_query = copy.copy(query)
-                shadow_query.async_future = None
-                shadow_query.metadata.is_shadow_query = True
-                for shadow_backend in shadow_backends:
-                    assigned_backends.add(shadow_backend)
-                    backend_queues[shadow_backend].appendleft(shadow_query)
-
-        return assigned_backends
+        chosen_backend, shadow_backends = self._select_backends(value)
+        logger.debug(f"Chosen backend {chosen_backend} for query {query}")
+        return [chosen_backend] + shadow_backends

python/ray/serve/handle.py

@@ -1,9 +1,27 @@
-from typing import Optional, Dict, Any, Union
+import asyncio
+import concurrent.futures
+import threading
+from typing import Any, Coroutine, Dict, Optional, Union
 
+import ray
 from ray.serve.context import TaskContext
-from ray.serve.router import RequestMetadata
+from ray.serve.router import RequestMetadata, Router
 from ray.serve.utils import get_random_letters
 
+global_async_loop = None
+
+
+def create_or_get_async_loop_in_thread():
+    global global_async_loop
+    if global_async_loop is None:
+        global_async_loop = asyncio.new_event_loop()
+        thread = threading.Thread(
+            daemon=True,
+            target=global_async_loop.run_forever,
+        )
+        thread.start()
+    return global_async_loop
+
 
 class RayServeHandle:
     """A handle to a service endpoint.
@@ -28,15 +46,16 @@ class RayServeHandle:
 
     def __init__(
             self,
-            router_handle,
+            controller_handle,
             endpoint_name,
+            sync: bool,
             *,
             method_name=None,
             shard_key=None,
             http_method=None,
             http_headers=None,
     ):
-        self.router_handle = router_handle
+        self.controller_handle = controller_handle
         self.endpoint_name = endpoint_name
 
         self.method_name = method_name
@@ -44,6 +63,36 @@ class RayServeHandle:
         self.http_method = http_method
         self.http_headers = http_headers
 
+        self.router = Router(self.controller_handle)
+        self.sync = sync
+        # In the synchrounous mode, we create a new event loop in a separate
+        # thread and run the Router.setup in that loop. In the async mode, we
+        # can just use the current loop we are in right now.
+        if self.sync:
+            self.async_loop = create_or_get_async_loop_in_thread()
+            asyncio.run_coroutine_threadsafe(
+                self.router.setup_in_async_loop(),
+                self.async_loop,
+            )
+        else:  # async
+            self.async_loop = asyncio.get_event_loop()
+            # create_task is not threadsafe.
+            self.async_loop.create_task(self.router.setup_in_async_loop())
+
+    def _remote(self, request_data, kwargs) -> Coroutine:
+        request_metadata = RequestMetadata(
+            get_random_letters(10),  # Used for debugging.
+            self.endpoint_name,
+            TaskContext.Python,
+            call_method=self.method_name or "__call__",
+            shard_key=self.shard_key,
+            http_method=self.http_method or "GET",
+            http_headers=self.http_headers or dict(),
+        )
+        coro = self.router.assign_request(request_metadata, request_data,
+                                          **kwargs)
+        return coro
+
     def remote(self, request_data: Optional[Union[Dict, Any]] = None,
                **kwargs):
         """Issue an asynchrounous request to the endpoint.
@@ -60,17 +109,17 @@ class RayServeHandle:
             ``**kwargs``: All keyword arguments will be available in
                 ``request.args``.
         """
-        request_metadata = RequestMetadata(
-            get_random_letters(10),  # Used for debugging.
-            self.endpoint_name,
-            TaskContext.Python,
-            call_method=self.method_name or "__call__",
-            shard_key=self.shard_key,
-            http_method=self.http_method or "GET",
-            http_headers=self.http_headers or dict(),
-        )
-        return self.router_handle.enqueue_request.remote(
-            request_metadata, request_data, **kwargs)
+        assert self.sync, "handle.remote() should be called from sync handle."
+        coro = self._remote(request_data, kwargs)
+        future: concurrent.futures.Future = asyncio.run_coroutine_threadsafe(
+            coro, self.async_loop)
+        # Block until the result is ready.
+        return future.result()
+
+    async def _remote_async(self, request_data, **kwargs) -> ray.ObjectRef:
+        """Experimental API for enqueue a request in async context."""
+        assert not self.sync, "_remote_async must be called inside async loop."
+        return await self._remote(request_data, kwargs)
 
     def options(self,
                 method_name: Optional[str] = None,
@@ -86,15 +135,12 @@ class RayServeHandle:
             shard_key(str): A string to use to deterministically map this
                 request to a backend if there are multiple for this endpoint.
         """
-        return RayServeHandle(
-            self.router_handle,
-            self.endpoint_name,
-            # Don't override existing method
-            method_name=self.method_name or method_name,
-            shard_key=self.shard_key or shard_key,
-            http_method=self.http_method or http_method,
-            http_headers=self.http_headers or http_headers,
-        )
+        # Don't override default non-null values.
+        self.method_name = self.method_name or method_name
+        self.shard_key = self.shard_key or shard_key
+        self.http_method = self.http_method or http_method
+        self.http_headers = self.http_headers or http_headers
+        return self
 
     def __repr__(self):
         return f"RayServeHandle(endpoint='{self.endpoint_name}')"

python/ray/serve/http_proxy.py

@@ -28,7 +28,7 @@ class HTTPProxy:
     # blocks forever
     """
 
-    async def fetch_config_from_controller(self, name, controller_name):
+    async def fetch_config_from_controller(self, controller_name):
         assert ray.is_initialized()
         controller = ray.get_actor(controller_name)
 
@@ -39,8 +39,8 @@ class HTTPProxy:
             description="The number of HTTP requests processed",
             tag_keys=("route", ))
 
-        self.router = Router()
-        await self.router.setup(name, controller_name)
+        self.router = Router(controller)
+        await self.router.setup_in_async_loop()
 
     def set_route_table(self, route_table):
         self.route_table = route_table
@@ -119,8 +119,9 @@ class HTTPProxy:
             shard_key=headers.get("X-SERVE-SHARD-KEY".lower(), None),
         )
 
-        result = await self.router.enqueue_request(request_metadata, scope,
-                                                   http_body_bytes)
+        ref = await self.router.assign_request(request_metadata, scope,
+                                               http_body_bytes)
+        result = await ref
 
         if isinstance(result, RayTaskError):
             error_message = "Task Error. Traceback: {}.".format(result)
@@ -133,17 +134,15 @@ class HTTPProxy:
 class HTTPProxyActor:
     async def __init__(
             self,
-            name,
             host,
             port,
             controller_name,
             http_middlewares: List["starlette.middleware.Middleware"] = []):
-        self.app = HTTPProxy()
         self.host = host
         self.port = port
 
         self.app = HTTPProxy()
-        await self.app.fetch_config_from_controller(name, controller_name)
+        await self.app.fetch_config_from_controller(controller_name)
 
         self.wrapped_app = self.app
         for middleware in http_middlewares:
@@ -186,7 +185,7 @@ class HTTPProxyActor:
         self.app.set_route_table(route_table)
 
     # ------ Proxy router logic ------ #
-    async def enqueue_request(self, request_meta, *request_args,
-                              **request_kwargs):
-        return await self.app.router.enqueue_request(
-            request_meta, *request_args, **request_kwargs)
+    async def assign_request(self, request_meta, *request_args,
+                             **request_kwargs):
+        return await (await self.app.router.assign_request(
+            request_meta, *request_args, **request_kwargs))

python/ray/serve/long_poll.py

@@ -65,6 +65,7 @@ class LongPollerAsyncClient:
         while True:
             updates: Dict[str, UpdatedObject] = await self._poll_once()
             self._update(updates)
+            logger.debug(f"LongPollerClient received udpates: {updates}")
             for key, updated_object in updates.items():
                 # NOTE(simon): This blocks the loop from doing another poll.
                 # Consider use loop.create_task here or poll first then call

python/ray/serve/router.py

@@ -1,18 +1,16 @@
 import asyncio
-import copy
-from collections import defaultdict, deque
-import time
-from typing import DefaultDict, List, Dict, Any, Optional
-import pickle
+import itertools
+from collections import defaultdict
 from dataclasses import dataclass, field
+from typing import Any, DefaultDict, Dict, Iterable, List, Optional
 
 import ray
-from ray.exceptions import RayTaskError
-from ray.serve.long_poll import LongPollerAsyncClient
-from ray.util import metrics
+from ray.actor import ActorHandle
 from ray.serve.context import TaskContext
-from ray.serve.endpoint_policy import RandomEndpointPolicy
-from ray.serve.utils import logger, chain_future
+from ray.serve.endpoint_policy import EndpointPolicy, RandomEndpointPolicy
+from ray.serve.long_poll import LongPollerAsyncClient
+from ray.serve.utils import logger
+from ray.util import metrics
 
 REPORT_QUEUE_LENGTH_PERIOD_S = 1.0
 
@@ -41,382 +39,200 @@ class Query:
     args: List[Any]
     kwargs: Dict[Any, Any]
     context: TaskContext
-
     metadata: RequestMetadata
-    async_future: Optional[asyncio.Future] = None
 
-    tick_enter_router: Optional[float] = None
+    # Fields used by backend worker to perform timing measurement.
     tick_enter_replica: Optional[float] = None
 
-    def __reduce__(self):
-        return type(self).ray_deserialize, (self.ray_serialize(), )
 
-    def ray_serialize(self):
-        # NOTE: this method is needed because Query need to be serialized and
-        # sent to the replica. However, after we send the query to the
-        # replica the async_future is still needed to retrieve the final
-        # result. Therefore we need a way to pass the information to replicas
-        # without removing async_future.
-        clone = copy.copy(self.__dict__)
-        clone.pop("async_future")
-        return pickle.dumps(clone)
+class ReplicaSet:
+    """Data structure representing a set of replica actor handles"""
 
-    @staticmethod
-    def ray_deserialize(value):
-        kwargs = pickle.loads(value)
-        return Query(**kwargs)
+    def __init__(self):
+        # NOTE(simon): We have to do this because max_concurrent_queries
+        # and the replica handles come from different long poll keys.
+        self.max_concurrent_queries: int = 8
+        self.in_flight_queries: Dict[ActorHandle, set] = dict()
+
+        # The iterator used for load balancing among replicas. Using itertools
+        # cycle, we implements a round-robin policy, skipping overloaded
+        # replicas.
+        # NOTE(simon): We can make this more pluggable and consider different
+        # policies like: min load, pick min of two replicas, pick replicas on
+        # the same node.
+        self.replica_iterator = itertools.cycle(self.in_flight_queries.keys())
+
+        # Used to unblock this replica set waiting for free replicas. A newly
+        # added replica or updated max_concurrenty_queries value means the
+        # query that waits on a free replica might be unblocked on.
+        self.config_updated_event = asyncio.Event()
+
+    def set_max_concurrent_queries(self, new_value):
+        if new_value != self.max_concurrent_queries:
+            self.max_concurrent_queries = new_value
+            logger.debug(
+                f"ReplicaSet: chaging max_concurrent_queries to {new_value}")
+            self.config_updated_event.set()
+
+    def update_worker_replicas(self, worker_replicas: Iterable[ActorHandle]):
+        current_replica_set = set(self.in_flight_queries.keys())
+        updated_replica_set = set(worker_replicas)
+
+        added = updated_replica_set - current_replica_set
+        for new_replica_handle in added:
+            self.in_flight_queries[new_replica_handle] = set()
+
+        removed = current_replica_set - updated_replica_set
+        for removed_replica_handle in removed:
+            # NOTE(simon): Do we warn if there are still inflight queries?
+            # The current approach is no because the queries objectrefs are
+            # just used to perform backpressure. Caller should decide what to
+            # do with the object refs.
+            del self.in_flight_queries[removed_replica_handle]
+
+        # State changed, reset the round robin iterator
+        if len(added) > 0 or len(removed) > 0:
+            self.replica_iterator = itertools.cycle(
+                self.in_flight_queries.keys())
+            self.config_updated_event.set()
+
+    def _try_assign_replica(self, query: Query) -> Optional[ray.ObjectRef]:
+        """Try to assign query to a replica, return the object ref is succeeded
+        or return None if it can't assign this query to any replicas.
+        """
+        for _ in range(len(self.in_flight_queries.keys())):
+            replica = next(self.replica_iterator)
+            if len(self.in_flight_queries[replica]
+                   ) >= self.max_concurrent_queries:
+                # This replica is overloaded, try next one
+                continue
+            logger.debug(f"Replica set assigned {query} to {replica}")
+            ref = replica.handle_request.remote(query)
+            self.in_flight_queries[replica].add(ref)
+            return ref
+        return None
+
+    @property
+    def _all_query_refs(self):
+        return list(
+            itertools.chain.from_iterable(self.in_flight_queries.values()))
+
+    def _drain_completed_object_refs(self) -> int:
+        refs = self._all_query_refs
+        done, _ = ray.wait(refs, num_returns=len(refs), timeout=0)
+        for replica_in_flight_queries in self.in_flight_queries.values():
+            replica_in_flight_queries.difference_update(done)
+        return len(done)
+
+    async def assign_replica(self, query: Query) -> ray.ObjectRef:
+        """Given a query, submit it to a replica and return the object ref.
+
+        This method will keep track of the in flight queries for each replicas
+        and only send a query to available replicas (determined by the backend
+        max_concurrent_quries value.)
+        """
+        assigned_ref = self._try_assign_replica(query)
+        while assigned_ref is None:  # Can't assign a replica right now.
+            logger.debug(f"Failed to assign a replica for query {query}")
+            # Maybe there exists a free replica, we just need to refresh our
+            # query tracker.
+            num_finished = self._drain_completed_object_refs()
+            # All replicas are really busy, wait for a query to complete or the
+            # config to be updated.
+            if num_finished == 0:
+                logger.debug(
+                    f"All replicas are busy, waiting for a free replica.")
+                await asyncio.wait(
+                    self._all_query_refs + [self.config_updated_event.wait()],
+                    return_when=asyncio.FIRST_COMPLETED)
+                if self.config_updated_event.is_set():
+                    self.config_updated_event.clear()
+            # We are pretty sure a free replica is ready now, let's recurse and
+            # assign this query a replica.
+            assigned_ref = await self.assign_replica(query)
+        return assigned_ref
 
 
 class Router:
-    """A router that routes request to available replicas."""
-
-    async def setup(self, name, controller_name, _do_long_pull=True):
-        """Setup the router state
+    def __init__(self, controller_handle: ActorHandle):
+        """Router process incoming queries: choose backend, and assign replica.
 
         Args:
-            name(str): Used to identify the router when reporting queue
-                lengths to the controller.
-            controller_name(str): The actor name for the controller.
-            _do_long_pull(bool): Used by unit testing.
+            controller_handle(ActorHandle): The controller handle.
         """
+        self.controller = controller_handle
 
-        # Note: Several queues are used in the router
-        # - When a request come in, it's placed inside its corresponding
-        #   endpoint_queue.
-        # - The endpoint_queue is dequeued during flush operation, which moves
-        #   the queries to backend buffer_queue. Here we match a request
-        #   for an endpoint to a backend given some policy.
-        # - The worker_queue is used to collect idle actor handle. These
-        #   handles are dequed during the second stage of flush operation,
-        #   which assign queries in buffer_queue to actor handle.
+        self.endpoint_policies: Dict[str, EndpointPolicy] = dict()
+        self.backend_replicas: Dict[str, ReplicaSet] = defaultdict(ReplicaSet)
 
-        self.name = name
-
-        # -- Queues -- #
-
-        # endpoint_name -> request queue
-        # We use FIFO (left to right) ordering. The new items should be added
-        # using appendleft. Old items should be removed via pop().
-        self.endpoint_queues: DefaultDict[deque[Query]] = defaultdict(deque)
-        # backend_name -> worker replica tag queue
-        self.worker_queues: DefaultDict[deque[str]] = defaultdict(deque)
-        # backend_name -> worker payload queue
-        self.backend_queues = defaultdict(deque)
-
-        # -- Metadata -- #
-
-        # endpoint_name -> traffic_policy
-        self.traffic = dict()
-        # backend_name -> backend_config
-        self.backend_info = dict()
-        # replica tag -> worker_handle
-        self.replicas = dict()
-        # backend_name -> replica_tag -> concurrent queries counter
-        self.queries_counter = defaultdict(lambda: defaultdict(int))
-
-        # -- Synchronization -- #
-
-        # This lock guarantee that only one flush operation can happen at a
-        # time. Without the lock, multiple flush operation can pop from the
-        # same buffer_queue and worker_queue and create deadlock. For example,
-        # an operation holding the only query and the other flush operation
-        # holding the only idle replica. Additionally, allowing only one flush
-        # operation at a time simplifies design overhead for custom queuing and
-        # batching policies.
-        self.flush_lock = asyncio.Lock()
-
-        # -- State Restoration -- #
-        # Fetch the replica handles, traffic policies, and backend configs from
-        # the controller. We use a "pull-based" approach instead of pushing
-        # them from the controller so that the router can transparently recover
-        # from failure.
-        self.controller = ray.get_actor(controller_name)
+        self._pending_endpoints: DefaultDict[str, asyncio.Event] = defaultdict(
+            asyncio.Event)
 
         # -- Metrics Registration -- #
         self.num_router_requests = metrics.Count(
             "num_router_requests",
             description="Number of requests processed by the router.",
             tag_keys=("endpoint", ))
-        self.num_error_endpoint_requests = metrics.Count(
-            "num_error_endpoint_requests",
-            description=(
-                "Number of requests that errored when getting results "
-                "for the endpoint."),
-            tag_keys=("endpoint", ))
-        self.num_error_backend_requests = metrics.Count(
-            "num_error_backend_requests",
-            description=("Number of requests that errored when getting result "
-                         "from the backend."),
-            tag_keys=("backend", ))
 
-        self.backend_queue_size = metrics.Gauge(
-            "backend_queued_queries",
-            description=("Current number of queries queued "
-                         "in the router for a backend"),
-            tag_keys=("backend", ))
+    async def setup_in_async_loop(self):
+        # NOTE(simon): Instead of performing initialization in __init__,
+        # We separated the init of LongPollerAsyncClient to this method because
+        # __init__ might be called in sync context. LongPollerAsyncClient
+        # requires async context.
+        self.long_pull_client = LongPollerAsyncClient(
+            self.controller, {
+                "traffic_policies": self._update_traffic_policies,
+                "worker_handles": self._update_worker_handles,
+                "backend_configs": self._update_backend_configs,
+            })
 
-        asyncio.get_event_loop().create_task(self.report_queue_lengths())
+    async def _update_traffic_policies(self, traffic_policies):
+        for endpoint, traffic_policy in traffic_policies.items():
+            self.endpoint_policies[endpoint] = RandomEndpointPolicy(
+                traffic_policy)
+            if endpoint in self._pending_endpoints:
+                event = self._pending_endpoints.pop(endpoint)
+                event.set()
 
-        if _do_long_pull:
-            self.long_poll_client = LongPollerAsyncClient(
-                self.controller, {
-                    "traffic_policies": self.update_traffic_policies,
-                    "worker_handles": self.update_worker_handles,
-                    "backend_configs": self.update_backend_configs
-                })
+    async def _update_worker_handles(self, worker_handles):
+        for backend_tag, replica_handles in worker_handles.items():
+            self.backend_replicas[backend_tag].update_worker_replicas(
+                replica_handles)
 
-    async def update_traffic_policies(self, traffic_policies):
-        updated_endpoints = set(traffic_policies.keys())
-        curr_endpoints = set(self.traffic.keys())
+    async def _update_backend_configs(self, backend_configs):
+        for backend_tag, config in backend_configs.items():
+            self.backend_replicas[backend_tag].set_max_concurrent_queries(
+                config.max_concurrent_queries)
 
-        for endpoint in updated_endpoints:
-            await self.set_traffic(endpoint, traffic_policies[endpoint])
-
-        removed_endpoints = curr_endpoints - updated_endpoints
-        for endpoint in removed_endpoints:
-            await self.remove_endpoint(endpoint)
-
-    async def update_worker_handles(self, worker_handles):
-        for backend_tag, replica_dict in worker_handles.items():
-            # NOTE(simon): This is a just hack around the current data
-            # structure to resolve replicas added and removed. It will be
-            # immediately become obselete when we update the router.
-            updated_replica_tags = set(replica_dict.keys())
-            curr_replica_tags = {
-                tag.replace(backend_tag + ":", "")
-                for tag in self.replicas.keys() if tag.startswith(backend_tag)
-            }
-
-            added_replicas = updated_replica_tags - curr_replica_tags
-            removed_replicas = curr_replica_tags - updated_replica_tags
-
-            for replica_tag in added_replicas:
-                await self.add_new_replica(backend_tag, replica_tag,
-                                           replica_dict[replica_tag])
-            for replica_tag in removed_replicas:
-                await self.remove_replica(backend_tag, replica_tag)
-
-    async def update_backend_configs(self, backend_configs):
-        updated_backends = set(backend_configs.keys())
-        curr_backends = set(self.backend_info.keys())
-
-        for backend in updated_backends:
-            await self.set_backend_config(backend, backend_configs[backend])
-
-        removed_backends = curr_backends - updated_backends
-        for backend in removed_backends:
-            await self.remove_backend(backend)
-
-    async def enqueue_request(self, request_meta, *request_args,
-                              **request_kwargs):
+    async def assign_request(
+            self,
+            request_meta: RequestMetadata,
+            *request_args,
+            **request_kwargs,
+    ):
+        """Assign a query and returns an object ref represent the result"""
         endpoint = request_meta.endpoint
-        logger.debug("Received request {} for endpoint {}.".format(
-            request_meta.request_id, endpoint))
-        request_start = time.time()
+        query = Query(
+            args=list(request_args),
+            kwargs=request_kwargs,
+            context=request_meta.request_context,
+            metadata=request_meta,
+        )
+
+        if endpoint not in self.endpoint_policies:
+            logger.info(
+                f"Endpoint {endpoint} doesn't exist, waiting for registration."
+            )
+            await self._pending_endpoints[endpoint].wait()
+
+        endpoint_policy = self.endpoint_policies[endpoint]
+        chosen_backend, *shadow_backends = endpoint_policy.assign(query)
+
+        result_ref = await self.backend_replicas[chosen_backend
+                                                 ].assign_replica(query)
+        for backend in shadow_backends:
+            await self.backend_replicas[backend].assign_replica(query)
+
         self.num_router_requests.record(1, tags={"endpoint": endpoint})
 
-        request_context = request_meta.request_context
-        query = Query(
-            request_args,
-            request_kwargs,
-            request_context,
-            metadata=request_meta,
-            async_future=asyncio.get_event_loop().create_future())
-        async with self.flush_lock:
-            self.endpoint_queues[endpoint].appendleft(query)
-            self.flush_endpoint_queue(endpoint)
-
-        try:
-            result = await query.async_future
-        except RayTaskError as e:
-            self.num_error_endpoint_requests.record(
-                1, tags={"endpoint": endpoint})
-            result = e
-
-        request_time_ms = (time.time() - request_start) * 1000
-        logger.debug("Finished request {} in {:.2f}ms".format(
-            request_meta.request_id, request_time_ms))
-        return result
-
-    async def add_new_replica(self, backend_tag, replica_tag, replica_handle):
-        backend_replica_tag = backend_tag + ":" + replica_tag
-        if backend_replica_tag in self.replicas:
-            return
-        self.replicas[backend_replica_tag] = replica_handle
-
-        logger.debug("New worker added for backend '{}'".format(backend_tag))
-        await self.mark_worker_idle(backend_tag, backend_replica_tag)
-
-    async def mark_worker_idle(self, backend_tag, backend_replica_tag):
-        logger.debug(
-            "Marking backend with tag {} as idle.".format(backend_replica_tag))
-        if backend_replica_tag not in self.replicas:
-            return
-
-        async with self.flush_lock:
-            # NOTE(simon): This is a O(n) operation where n=len(worker_queue)
-            if backend_replica_tag not in self.worker_queues[backend_tag]:
-                self.worker_queues[backend_tag].appendleft(backend_replica_tag)
-            self.flush_backend_queues([backend_tag])
-
-    async def remove_replica(self, backend_tag, replica_tag):
-        backend_replica_tag = backend_tag + ":" + replica_tag
-        if backend_replica_tag not in self.replicas:
-            return
-
-        # We need this lock because we modify worker_queue here.
-        async with self.flush_lock:
-            del self.replicas[backend_replica_tag]
-
-            try:
-                self.worker_queues[backend_tag].remove(backend_replica_tag)
-            except ValueError:
-                # Replica doesn't exist in the idle worker queues.
-                # It's ok because the worker might not have returned the
-                # result.
-                pass
-
-    async def set_traffic(self, endpoint, traffic_policy):
-        logger.debug("Setting traffic for endpoint %s to %s", endpoint,
-                     traffic_policy)
-        async with self.flush_lock:
-            self.traffic[endpoint] = RandomEndpointPolicy(traffic_policy)
-            self.flush_endpoint_queue(endpoint)
-
-    async def remove_endpoint(self, endpoint):
-        logger.debug("Removing endpoint {}".format(endpoint))
-        async with self.flush_lock:
-            self.flush_endpoint_queue(endpoint)
-            if endpoint in self.endpoint_queues:
-                del self.endpoint_queues[endpoint]
-            if endpoint in self.traffic:
-                del self.traffic[endpoint]
-
-    async def set_backend_config(self, backend, config):
-        logger.debug("Setting backend config for "
-                     "backend {} to {}.".format(backend, config))
-        async with self.flush_lock:
-            self.backend_info[backend] = config
-
-    async def remove_backend(self, backend):
-        logger.debug("Removing backend {}".format(backend))
-        async with self.flush_lock:
-            self.flush_backend_queues([backend])
-            if backend in self.backend_info:
-                del self.backend_info[backend]
-            if backend in self.worker_queues:
-                del self.worker_queues[backend]
-            if backend in self.backend_queues:
-                del self.backend_queues[backend]
-
-    def flush_endpoint_queue(self, endpoint):
-        """Attempt to schedule any pending requests to available backends."""
-        assert self.flush_lock.locked()
-        if endpoint not in self.traffic:
-            return
-        backends_to_flush = self.traffic[endpoint].flush(
-            self.endpoint_queues[endpoint], self.backend_queues)
-        self.flush_backend_queues(backends_to_flush)
-
-    # Flushes the specified backend queues and assigns work to workers.
-    def flush_backend_queues(self, backends_to_flush):
-        assert self.flush_lock.locked()
-        for backend in backends_to_flush:
-            # No workers available.
-            if len(self.worker_queues[backend]) == 0:
-                continue
-            # No work to do.
-            if len(self.backend_queues[backend]) == 0:
-                continue
-
-            buffer_queue = self.backend_queues[backend]
-            worker_queue = self.worker_queues[backend]
-
-            logger.debug("Assigning queries for backend {} with buffer "
-                         "queue size {} and worker queue size {}".format(
-                             backend, len(buffer_queue), len(worker_queue)))
-
-            self._assign_query_to_worker(
-                backend,
-                buffer_queue,
-                worker_queue,
-            )
-
-    async def _do_query(self, backend, backend_replica_tag, req):
-        # If the worker died, this will be a RayActorError. Just return it and
-        # let the HTTP proxy handle the retry logic.
-        logger.debug("Sending query to replica:" + backend_replica_tag)
-        worker = self.replicas[backend_replica_tag]
-        try:
-            object_ref = worker.handle_request.remote(req.ray_serialize())
-            if req.metadata.is_shadow_query:
-                # No need to actually get the result, but we do need to wait
-                # until the call completes to mark the worker idle.
-                await asyncio.wait([object_ref])
-                result = ""
-            else:
-                result = await object_ref
-        except RayTaskError as error:
-            self.num_error_backend_requests.record(
-                1, tags={"backend": backend})
-            result = error
-        self.queries_counter[backend][backend_replica_tag] -= 1
-        await self.mark_worker_idle(backend, backend_replica_tag)
-        return result
-
-    def _assign_query_to_worker(self, backend, buffer_queue, worker_queue):
-        overloaded_replicas = set()
-        while len(buffer_queue) and len(worker_queue):
-            backend_replica_tag = worker_queue.pop()
-
-            # The replica might have been deleted already.
-            if backend_replica_tag not in self.replicas:
-                continue
-
-            # We have reached the end of the worker queue where all replicas
-            # are overloaded.
-            if backend_replica_tag in overloaded_replicas:
-                break
-
-            # This replica has too many in flight and processing queries.
-            max_queries = 1
-            if backend in self.backend_info:
-                max_queries = self.backend_info[backend].max_concurrent_queries
-            curr_queries = self.queries_counter[backend][backend_replica_tag]
-            if curr_queries >= max_queries:
-                # Put the worker back to the queue.
-                worker_queue.appendleft(backend_replica_tag)
-                overloaded_replicas.add(backend_replica_tag)
-                logger.debug(
-                    "Skipping backend {} because it has {} in flight "
-                    "requests which exceeded the concurrency limit.".format(
-                        backend, curr_queries))
-                continue
-
-            request = buffer_queue.pop()
-            logger.debug("Assigning request {} to replica {}.".format(
-                request.metadata.request_id, backend_replica_tag))
-            self.queries_counter[backend][backend_replica_tag] += 1
-            future = asyncio.get_event_loop().create_task(
-                self._do_query(backend, backend_replica_tag, request))
-
-            # For shadow queries, just ignore the result.
-            if not request.metadata.is_shadow_query:
-                chain_future(future, request.async_future)
-
-            worker_queue.appendleft(backend_replica_tag)
-
-    async def report_queue_lengths(self):
-        while True:
-            queue_lengths = {
-                backend: len(q)
-                for backend, q in self.backend_queues.items()
-            }
-            self.controller.report_queue_lengths.remote(
-                self.name, queue_lengths)
-
-            for backend, length in queue_lengths.items():
-                self.backend_queue_size.record(
-                    length, tags={"backend": backend})
-
-            await asyncio.sleep(REPORT_QUEUE_LENGTH_PERIOD_S)
+        return result_ref

python/ray/serve/tests/test_backend_worker.py

@@ -47,14 +47,17 @@ def setup_worker(name,
 
 async def add_servable_to_router(servable, router, **kwargs):
     worker = setup_worker("backend", servable, **kwargs)
-    await router.add_new_replica.remote("backend", "replica", worker)
-    await router.set_traffic.remote("endpoint", TrafficPolicy({
-        "backend": 1.0
-    }))
+    await router._update_worker_handles.remote({"backend": [worker]})
+    await router._update_traffic_policies.remote({
+        "endpoint": TrafficPolicy({
+            "backend": 1.0
+        })
+    })
 
     if "backend_config" in kwargs:
-        await router.set_backend_config.remote("backend",
-                                               kwargs["backend_config"])
+        await router._update_backend_configs.remote({
+            "backend": kwargs["backend_config"]
+        })
     return worker
 
 
@@ -67,9 +70,8 @@ def make_request_param(call_method="__call__"):
 
 
 @pytest.fixture
-def router(serve_instance):
-    q = ray.remote(Router).remote()
-    ray.get(q.setup.remote("", serve_instance._controller_name))
+async def router(serve_instance):
+    q = ray.remote(Router).remote(serve_instance._controller)
     yield q
     ray.kill(q)
 
@@ -87,7 +89,8 @@ async def test_servable_function(serve_instance, router):
 
     for query in [333, 444, 555]:
         query_param = make_request_param()
-        result = await router.enqueue_request.remote(query_param, i=query)
+        result = await (await router.assign_request.remote(
+            query_param, i=query))
         assert result == query
 
 
@@ -103,7 +106,8 @@ async def test_servable_class(serve_instance, router):
 
     for query in [333, 444, 555]:
         query_param = make_request_param()
-        result = await router.enqueue_request.remote(query_param, i=query)
+        result = await (await router.assign_request.remote(
+            query_param, i=query))
         assert result == query + 3
 
 
@@ -118,16 +122,16 @@ async def test_task_runner_custom_method_single(serve_instance, router):
     _ = await add_servable_to_router(NonBatcher, router)
 
     query_param = make_request_param("a")
-    a_result = await router.enqueue_request.remote(query_param)
+    a_result = await (await router.assign_request.remote(query_param))
     assert a_result == "a"
 
     query_param = make_request_param("b")
-    b_result = await router.enqueue_request.remote(query_param)
+    b_result = await (await router.assign_request.remote(query_param))
     assert b_result == "b"
 
     query_param = make_request_param("non_exist")
     with pytest.raises(ray.exceptions.RayTaskError):
-        await router.enqueue_request.remote(query_param)
+        await (await router.assign_request.remote(query_param))
 
 
 async def test_task_runner_custom_method_batch(serve_instance, router):
@@ -149,8 +153,12 @@ async def test_task_runner_custom_method_batch(serve_instance, router):
     a_query_param = make_request_param("a")
     b_query_param = make_request_param("b")
 
-    futures = [router.enqueue_request.remote(a_query_param) for _ in range(2)]
-    futures += [router.enqueue_request.remote(b_query_param) for _ in range(2)]
+    futures = [
+        await router.assign_request.remote(a_query_param) for _ in range(2)
+    ]
+    futures += [
+        await router.assign_request.remote(b_query_param) for _ in range(2)
+    ]
 
     gathered = await asyncio.gather(*futures)
     assert set(gathered) == {"a-0", "a-1", "b-0", "b-1"}
@@ -176,14 +184,14 @@ async def test_servable_batch_error(serve_instance, router):
 
     with pytest.raises(RayServeException, match="doesn't preserve batch size"):
         different_size = make_request_param("error_different_size")
-        await router.enqueue_request.remote(different_size)
+        await (await router.assign_request.remote(different_size))
 
     with pytest.raises(RayServeException, match="iterable"):
         non_iterable = make_request_param("error_non_iterable")
-        await router.enqueue_request.remote(non_iterable)
+        await (await router.assign_request.remote(non_iterable))
 
     np_array = make_request_param("return_np_array")
-    result_np_value = await router.enqueue_request.remote(np_array)
+    result_np_value = await (await router.assign_request.remote(np_array))
     assert isinstance(result_np_value, np.int32)
 
 
@@ -200,8 +208,8 @@ async def test_task_runner_perform_batch(serve_instance, router):
     _ = await add_servable_to_router(batcher, router, backend_config=config)
 
     query_param = make_request_param()
-    my_batch_sizes = await asyncio.gather(
-        *[router.enqueue_request.remote(query_param) for _ in range(3)])
+    my_batch_sizes = await asyncio.gather(*[(
+        await router.assign_request.remote(query_param)) for _ in range(3)])
     assert my_batch_sizes == [2, 2, 1]
 
 
@@ -236,7 +244,7 @@ async def test_task_runner_perform_async(serve_instance, router):
     query_param = make_request_param()
 
     done, not_done = await asyncio.wait(
-        [router.enqueue_request.remote(query_param) for _ in range(10)],
+        [(await router.assign_request.remote(query_param)) for _ in range(10)],
         timeout=10)
     assert len(done) == 10
     for item in done:

python/ray/serve/tests/test_router.py

@@ -2,21 +2,31 @@
 Unit tests for the router class. Please don't add any test that will involve
 controller or the backend worker, use mock if necessary.
 """
-
+import asyncio
 from collections import defaultdict
+import os
 
 import pytest
-import ray
+from ray.serve.context import TaskContext
 
+import ray
+from ray.serve.config import BackendConfig
 from ray.serve.controller import TrafficPolicy
-from ray.serve.router import Router, Query, RequestMetadata
+from ray.serve.router import Query, ReplicaSet, RequestMetadata, Router
 from ray.serve.utils import get_random_letters
 from ray.test_utils import SignalActor
-from ray.serve.config import BackendConfig
 
 pytestmark = pytest.mark.asyncio
 
 
+@pytest.fixture
+def ray_instance():
+    os.environ["SERVE_LOG_DEBUG"] = "1"  # Turns on debug log for tests
+    ray.init(num_cpus=16)
+    yield
+    ray.shutdown()
+
+
 def mock_task_runner():
     @ray.remote(num_cpus=0)
     class TaskRunnerMock:
@@ -51,18 +61,62 @@ def task_runner_mock_actor():
     yield mock_task_runner()
 
 
-async def test_single_prod_cons_queue(serve_instance, task_runner_mock_actor):
-    q = ray.remote(Router).remote()
-    await q.setup.remote(
-        "", serve_instance._controller_name, _do_long_pull=False)
+@pytest.fixture
+def mock_controller():
+    @ray.remote(num_cpus=0)
+    class MockControllerActor:
+        def __init__(self):
+            from ray.serve.long_poll import LongPollerHost
+            self.host = LongPollerHost()
+            self.backend_replicas = defaultdict(list)
+            self.backend_configs = dict()
+            self.clear()
 
-    q.set_traffic.remote("svc", TrafficPolicy({"backend-single-prod": 1.0}))
-    q.add_new_replica.remote("backend-single-prod", "replica-1",
-                             task_runner_mock_actor)
+        def clear(self):
+            self.host.notify_changed("worker_handles", {})
+            self.host.notify_changed("traffic_policies", {})
+            self.host.notify_changed("backend_configs", {})
+
+        async def listen_for_change(self, snapshot_ids):
+            return await self.host.listen_for_change(snapshot_ids)
+
+        def set_traffic(self, endpoint, traffic_policy):
+            self.host.notify_changed("traffic_policies",
+                                     {endpoint: traffic_policy})
+
+        def add_new_replica(self,
+                            backend_tag,
+                            runner_actor,
+                            backend_config=BackendConfig()):
+            self.backend_replicas[backend_tag].append(runner_actor)
+            self.backend_configs[backend_tag] = backend_config
+
+            self.host.notify_changed(
+                "worker_handles",
+                self.backend_replicas,
+            )
+            self.host.notify_changed("backend_configs", self.backend_configs)
+
+    yield MockControllerActor.remote()
+
+
+async def test_simple_endpoint_backend_pair(ray_instance, mock_controller,
+                                            task_runner_mock_actor):
+    q = ray.remote(Router).remote(mock_controller)
+    await q.setup_in_async_loop.remote()
+
+    # Propogate configs
+    await mock_controller.set_traffic.remote(
+        "svc", TrafficPolicy({
+            "backend-single-prod": 1.0
+        }))
+    await mock_controller.add_new_replica.remote("backend-single-prod",
+                                                 task_runner_mock_actor)
 
     # Make sure we get the request result back
-    result = await q.enqueue_request.remote(
+    ref = await q.assign_request.remote(
         RequestMetadata(get_random_letters(10), "svc", None), 1)
+    result = await ref
     assert result == "DONE"
 
     # Make sure it's the right request
@@ -71,46 +125,53 @@ async def test_single_prod_cons_queue(serve_instance, task_runner_mock_actor):
     assert got_work.kwargs == {}
 
 
-async def test_alter_backend(serve_instance, task_runner_mock_actor):
-    q = ray.remote(Router).remote()
-    await q.setup.remote(
-        "", serve_instance._controller_name, _do_long_pull=False)
+async def test_changing_backend(ray_instance, mock_controller,
+                                task_runner_mock_actor):
+    q = ray.remote(Router).remote(mock_controller)
+    await q.setup_in_async_loop.remote()
 
-    await q.set_traffic.remote("svc", TrafficPolicy({"backend-alter": 1}))
-    await q.add_new_replica.remote("backend-alter", "replica-1",
-                                   task_runner_mock_actor)
-    await q.enqueue_request.remote(
+    await mock_controller.set_traffic.remote(
+        "svc", TrafficPolicy({
+            "backend-alter": 1
+        }))
+    await mock_controller.add_new_replica.remote("backend-alter",
+                                                 task_runner_mock_actor)
+
+    await q.assign_request.remote(
         RequestMetadata(get_random_letters(10), "svc", None), 1)
     got_work = await task_runner_mock_actor.get_recent_call.remote()
     assert got_work.args[0] == 1
 
-    await q.set_traffic.remote("svc", TrafficPolicy({"backend-alter-2": 1}))
-    await q.add_new_replica.remote("backend-alter-2", "replica-1",
-                                   task_runner_mock_actor)
-    await q.enqueue_request.remote(
+    await mock_controller.set_traffic.remote(
+        "svc", TrafficPolicy({
+            "backend-alter-2": 1
+        }))
+    await mock_controller.add_new_replica.remote("backend-alter-2",
+                                                 task_runner_mock_actor)
+    await q.assign_request.remote(
         RequestMetadata(get_random_letters(10), "svc", None), 2)
     got_work = await task_runner_mock_actor.get_recent_call.remote()
     assert got_work.args[0] == 2
 
 
-async def test_split_traffic_random(serve_instance, task_runner_mock_actor):
-    q = ray.remote(Router).remote()
-    await q.setup.remote(
-        "", serve_instance._controller_name, _do_long_pull=False)
+async def test_split_traffic_random(ray_instance, mock_controller,
+                                    task_runner_mock_actor):
+    q = ray.remote(Router).remote(mock_controller)
+    await q.setup_in_async_loop.remote()
 
-    await q.set_traffic.remote(
+    await mock_controller.set_traffic.remote(
         "svc", TrafficPolicy({
             "backend-split": 0.5,
             "backend-split-2": 0.5
         }))
     runner_1, runner_2 = [mock_task_runner() for _ in range(2)]
-    await q.add_new_replica.remote("backend-split", "replica-1", runner_1)
-    await q.add_new_replica.remote("backend-split-2", "replica-1", runner_2)
+    await mock_controller.add_new_replica.remote("backend-split", runner_1)
+    await mock_controller.add_new_replica.remote("backend-split-2", runner_2)
 
     # assume 50% split, the probability of all 20 requests goes to a
     # single queue is 0.5^20 ~ 1-6
     for _ in range(20):
-        await q.enqueue_request.remote(
+        await q.assign_request.remote(
             RequestMetadata(get_random_letters(10), "svc", None), 1)
 
     got_work = [
@@ -120,24 +181,10 @@ async def test_split_traffic_random(serve_instance, task_runner_mock_actor):
     assert [g.args[0] for g in got_work] == [1, 1]
 
 
-async def test_queue_remove_replicas(serve_instance):
-    class TestRouter(Router):
-        def worker_queue_size(self, backend):
-            return len(self.worker_queues["backend-remove"])
-
-    temp_actor = mock_task_runner()
-    q = ray.remote(TestRouter).remote()
-    await q.setup.remote(
-        "", serve_instance._controller_name, _do_long_pull=False)
-    await q.add_new_replica.remote("backend-remove", "replica-1", temp_actor)
-    await q.remove_replica.remote("backend-remove", "replica-1")
-    assert ray.get(q.worker_queue_size.remote("backend")) == 0
-
-
-async def test_shard_key(serve_instance, task_runner_mock_actor):
-    q = ray.remote(Router).remote()
-    await q.setup.remote(
-        "", serve_instance._controller_name, _do_long_pull=False)
+async def test_shard_key(ray_instance, mock_controller,
+                         task_runner_mock_actor):
+    q = ray.remote(Router).remote(mock_controller)
+    await q.setup_in_async_loop.remote()
 
     num_backends = 5
     traffic_dict = {}
@@ -145,13 +192,14 @@ async def test_shard_key(serve_instance, task_runner_mock_actor):
     for i, runner in enumerate(runners):
         backend_name = "backend-split-" + str(i)
         traffic_dict[backend_name] = 1.0 / num_backends
-        await q.add_new_replica.remote(backend_name, "replica-1", runner)
-    await q.set_traffic.remote("svc", TrafficPolicy(traffic_dict))
+        await mock_controller.add_new_replica.remote(backend_name, runner)
+    await mock_controller.set_traffic.remote("svc",
+                                             TrafficPolicy(traffic_dict))
 
     # Generate random shard keys and send one request for each.
     shard_keys = [get_random_letters() for _ in range(100)]
     for shard_key in shard_keys:
-        await q.enqueue_request.remote(
+        await q.assign_request.remote(
             RequestMetadata(
                 get_random_letters(10), "svc", None, shard_key=shard_key),
             shard_key)
@@ -166,7 +214,7 @@ async def test_shard_key(serve_instance, task_runner_mock_actor):
 
     # Send queries with the same shard keys a second time.
     for shard_key in shard_keys:
-        await q.enqueue_request.remote(
+        await q.assign_request.remote(
             RequestMetadata(
                 get_random_letters(10), "svc", None, shard_key=shard_key),
             shard_key)
@@ -178,71 +226,70 @@ async def test_shard_key(serve_instance, task_runner_mock_actor):
             assert call.args[0] in runner_shard_keys[i]
 
 
-async def test_router_use_max_concurrency(serve_instance):
+async def test_replica_set(ray_instance):
     signal = SignalActor.remote()
 
-    @ray.remote
+    @ray.remote(num_cpus=0)
     class MockWorker:
+        _num_queries = 0
+
         async def handle_request(self, request):
+            self._num_queries += 1
             await signal.wait.remote()
             return "DONE"
 
-        def ready(self):
-            pass
+        async def num_queries(self):
+            return self._num_queries
 
-    class VisibleRouter(Router):
-        def get_queues(self):
-            return self.queries_counter, self.backend_queues
+    # We will test a scenario with two replicas in the replica set.
+    rs = ReplicaSet()
+    workers = [MockWorker.remote() for _ in range(2)]
+    rs.set_max_concurrent_queries(1)
+    rs.update_worker_replicas(workers)
 
-    worker = MockWorker.remote()
-    q = ray.remote(VisibleRouter).remote()
-    await q.setup.remote(
-        "", serve_instance._controller_name, _do_long_pull=False)
-    backend_name = "max-concurrent-test"
-    config = BackendConfig(max_concurrent_queries=1)
-    await q.set_traffic.remote("svc", TrafficPolicy({backend_name: 1.0}))
-    await q.add_new_replica.remote(backend_name, "replica-tag", worker)
-    await q.set_backend_config.remote(backend_name, config)
+    # Send two queries. They should go through the router but blocked by signal
+    # actors.
+    query = Query([], {}, TaskContext.Python,
+                  RequestMetadata("request-id", "endpoint",
+                                  TaskContext.Python))
+    first_ref = await rs.assign_replica(query)
+    second_ref = await rs.assign_replica(query)
 
-    # We send over two queries
-    first_query = q.enqueue_request.remote(
-        RequestMetadata(get_random_letters(10), "svc", None), 1)
-    second_query = q.enqueue_request.remote(
-        RequestMetadata(get_random_letters(10), "svc", None), 1)
-
-    # Neither queries should be available
+    # These should be blocked by signal actor.
     with pytest.raises(ray.exceptions.GetTimeoutError):
-        ray.get([first_query, second_query], timeout=0.2)
+        ray.get([first_ref, second_ref], timeout=1)
 
-    # Let's retrieve the router internal state
-    queries_counter, backend_queues = await q.get_queues.remote()
-    # There should be just one inflight request
-    assert queries_counter[backend_name][
-        "max-concurrent-test:replica-tag"] == 1
-    # The second query is buffered
-    assert len(backend_queues["max-concurrent-test"]) == 1
+    # Each replica should have exactly one inflight query. Let make sure the
+    # queries arrived there.
+    for worker in workers:
+        while await worker.num_queries.remote() != 1:
+            await asyncio.sleep(1)
 
-    # Let's unblock the first query
-    await signal.send.remote(clear=True)
-    assert await first_query == "DONE"
+    # Let's try to send another query.
+    third_ref_pending_task = asyncio.get_event_loop().create_task(
+        rs.assign_replica(query))
+    # We should fail to assign a replica, so this coroutine should still be
+    # pending after some time.
+    await asyncio.sleep(0.2)
+    assert not third_ref_pending_task.done()
 
-    # The internal state of router should have changed.
-    queries_counter, backend_queues = await q.get_queues.remote()
-    # There should still be one inflight request
-    assert queries_counter[backend_name][
-        "max-concurrent-test:replica-tag"] == 1
-    # But there shouldn't be any queries in the queue
-    assert len(backend_queues["max-concurrent-test"]) == 0
+    # Let's unblock the two workers
+    await signal.send.remote()
+    assert await first_ref == "DONE"
+    assert await second_ref == "DONE"
 
-    # Unblocking the second query
-    await signal.send.remote(clear=True)
-    assert await second_query == "DONE"
+    # The third request should be unblocked and sent to first worker.
+    # This meas we should be able to get the object ref.
+    third_ref = await third_ref_pending_task
 
-    # Checking the internal state of the router one more time
-    queries_counter, backend_queues = await q.get_queues.remote()
-    assert queries_counter[backend_name][
-        "max-concurrent-test:replica-tag"] == 0
-    assert len(backend_queues["max-concurrent-test"]) == 0
+    # Now we got the object ref, let's get it result.
+    await signal.send.remote()
+    assert await third_ref == "DONE"
+
+    # Finally, make sure that one of the replica processed the third query.
+    num_queries_set = {(await worker.num_queries.remote())
+                       for worker in workers}
+    assert num_queries_set == {2, 1}
 
 
 if __name__ == "__main__":

python/ray/serve/tests/test_standalone.py

@@ -2,7 +2,6 @@
 The test file for all standalone tests that doesn't
 requires a shared Serve instance.
 """
-from random import randint
 import sys
 import socket
 
@@ -14,7 +13,7 @@ from ray import serve
 from ray.cluster_utils import Cluster
 from ray.serve.constants import SERVE_PROXY_NAME
 from ray.serve.utils import (block_until_http_ready, get_all_node_ids,
-                             format_actor_name, get_node_id_for_actor)
+                             format_actor_name)
 from ray.test_utils import wait_for_condition
 from ray._private.services import new_port
 
@@ -160,48 +159,5 @@ def test_middleware():
     ray.shutdown()
 
 
-@pytest.mark.skipif(
-    not hasattr(socket, "SO_REUSEPORT"),
-    reason=("Port sharing only works on newer verion of Linux. "
-            "This test can only be ran when port sharing is supported."))
-def test_cluster_handle_affinity():
-    cluster = Cluster()
-    # HACK: using two different ip address so the placement constraint for
-    # resource check later will work.
-    head_node = cluster.add_node(node_ip_address="127.0.0.1", num_cpus=4)
-    cluster.add_node(node_ip_address="0.0.0.0", num_cpus=4)
-
-    ray.init(head_node.address)
-
-    # Make sure we have two nodes.
-    node_ids = [n["NodeID"] for n in ray.nodes()]
-    assert len(node_ids) == 2
-
-    # Start the backend.
-    client = serve.start(http_port=randint(10000, 30000), detached=True)
-    client.create_backend("hi:v0", lambda _: "hi")
-    client.create_endpoint("hi", backend="hi:v0")
-
-    # Try to retrieve the handle from both head and worker node, check the
-    # router's node id.
-    @ray.remote
-    def check_handle_router_id():
-        client = serve.connect()
-        handle = client.get_handle("hi")
-        return get_node_id_for_actor(handle.router_handle)
-
-    router_node_ids = ray.get([
-        check_handle_router_id.options(resources={
-            node_id: 0.01
-        }).remote() for node_id in ray.state.node_ids()
-    ])
-
-    assert set(router_node_ids) == set(node_ids)
-
-    # Clean up the nodes (otherwise Ray will segfault).
-    ray.shutdown()
-    cluster.shutdown()
-
-
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", "-s", __file__]))