diff --git a/python/ray/serve/api.py b/python/ray/serve/api.py index 864039881..1e6a665f3 100644 --- a/python/ray/serve/api.py +++ b/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, diff --git a/python/ray/serve/controller.py b/python/ray/serve/controller.py index caac3a98a..cc702f4e1 100644 --- a/python/ray/serve/controller.py +++ b/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"" + 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( diff --git a/python/ray/serve/endpoint_policy.py b/python/ray/serve/endpoint_policy.py index 86728e68f..555c4f58a 100644 --- a/python/ray/serve/endpoint_policy.py +++ b/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 diff --git a/python/ray/serve/handle.py b/python/ray/serve/handle.py index 5d6d5616d..adcd5ac53 100644 --- a/python/ray/serve/handle.py +++ b/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}')" diff --git a/python/ray/serve/http_proxy.py b/python/ray/serve/http_proxy.py index da0766391..9432deefb 100644 --- a/python/ray/serve/http_proxy.py +++ b/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)) diff --git a/python/ray/serve/long_poll.py b/python/ray/serve/long_poll.py index 7c389873d..b178e1ba1 100644 --- a/python/ray/serve/long_poll.py +++ b/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 diff --git a/python/ray/serve/router.py b/python/ray/serve/router.py index 6573ad214..9808ad9c9 100644 --- a/python/ray/serve/router.py +++ b/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 diff --git a/python/ray/serve/tests/test_backend_worker.py b/python/ray/serve/tests/test_backend_worker.py index cdd7e8fd3..423c0db0e 100644 --- a/python/ray/serve/tests/test_backend_worker.py +++ b/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: diff --git a/python/ray/serve/tests/test_router.py b/python/ray/serve/tests/test_router.py index 08fbbec39..4752cc399 100644 --- a/python/ray/serve/tests/test_router.py +++ b/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__": diff --git a/python/ray/serve/tests/test_standalone.py b/python/ray/serve/tests/test_standalone.py index a32518028..686571068 100644 --- a/python/ray/serve/tests/test_standalone.py +++ b/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__]))