[serve] Refactor router policies to remove inheritance (#8372)

This commit is contained in:
Edward Oakes
2020-05-12 12:34:01 -05:00
committed by GitHub
parent 7ce138a6dc
commit bb494a3be8
5 changed files with 219 additions and 246 deletions
+48 -57
View File
@@ -103,7 +103,7 @@ class Router:
3. When there is only 1 backend ready, we will only use that backend.
"""
async def __init__(self):
async def __init__(self, policy, policy_kwargs):
# Note: Several queues are used in the router
# - When a request come in, it's placed inside its corresponding
# endpoint_queue.
@@ -114,6 +114,11 @@ class Router:
# handles are dequed during the second stage of flush operation,
# which assign queries in buffer_queue to actor handle.
# policy.RoutePolicy.
self.policy = policy
# kwargs to pass into the policy when it's constructed.
self.policy_kwargs = policy_kwargs
# -- Queues -- #
# endpoint_name -> request queue
@@ -123,12 +128,12 @@ class Router:
self.worker_queues: DefaultDict[asyncio.Queue[
ray.actor.ActorHandle]] = defaultdict(asyncio.Queue)
# backend_name -> worker payload queue
self.buffer_queues = defaultdict(blist.sortedlist)
self.backend_queues = defaultdict(blist.sortedlist)
# -- Metadata -- #
# endpoint_name -> traffic_policy
self.traffic = defaultdict(dict)
self.traffic = dict()
# backend_name -> backend_config
self.backend_info = dict()
# replica tag -> worker_handle
@@ -208,7 +213,8 @@ class Router:
call_method=request_meta.call_method,
async_future=asyncio.get_event_loop().create_future())
await self.endpoint_queues[endpoint].put(query)
await self.flush()
async with self.flush_lock:
await self.flush_endpoint_queue(endpoint)
# Note: a future change can be to directly return the ObjectID from
# replica task submission
@@ -226,7 +232,6 @@ class Router:
self.replicas[backend_replica_tag] = worker_handle
logger.debug("New worker added for backend '{}'".format(backend_tag))
# await worker_handle.ready.remote()
await self.mark_worker_idle(backend_tag, backend_replica_tag)
async def mark_worker_idle(self, backend_tag, backend_replica_tag):
@@ -234,7 +239,8 @@ class Router:
return
await self.worker_queues[backend_tag].put(backend_replica_tag)
await self.flush()
async with self.flush_lock:
await self.flush_backend_queues([backend_tag])
async def remove_worker(self, backend_tag, replica_tag):
backend_replica_tag = backend_tag + ":" + replica_tag
@@ -261,14 +267,15 @@ class Router:
async def set_traffic(self, endpoint, traffic_dict):
logger.debug("Setting traffic for endpoint %s to %s", endpoint,
traffic_dict)
self.traffic[endpoint] = traffic_dict
await self.flush()
async with self.flush_lock:
self.traffic[endpoint] = self.policy(traffic_dict,
**self.policy_kwargs)
await self.flush_endpoint_queue(endpoint)
async def remove_endpoint(self, endpoint):
logger.debug("Removing endpoint {}".format(endpoint))
async with self.flush_lock:
await self._flush_endpoint_queues()
await self._flush_buffer_queues()
await self.flush_endpoint_queue(endpoint)
if endpoint in self.endpoint_queues:
del self.endpoint_queues[endpoint]
if endpoint in self.traffic:
@@ -282,24 +289,22 @@ class Router:
async def remove_backend(self, backend):
logger.debug("Removing backend {}".format(backend))
async with self.flush_lock:
await self._flush_endpoint_queues()
await self._flush_buffer_queues()
await 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.buffer_queues:
del self.buffer_queues[backend]
if backend in self.backend_queues:
del self.backend_queues[backend]
async def flush(self):
"""In the default case, flush calls ._flush.
When this class is a Ray actor, .flush can be scheduled as a remote
method invocation.
"""
async with self.flush_lock:
await self._flush_endpoint_queues()
await self._flush_buffer_queues()
async 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 = await self.traffic[endpoint].flush(
self.endpoint_queues[endpoint], self.backend_queues)
await self.flush_backend_queues(backends_to_flush)
def _get_available_backends(self, endpoint):
backends_in_policy = set(self.traffic[endpoint].keys())
@@ -310,44 +315,30 @@ class Router:
}
return list(backends_in_policy.intersection(available_workers))
async def _flush_endpoint_queues(self):
"""Selects the backend and puts the endpoint queue query to the buffer
Expected Implementation:
The implementer is expected to access and manipulate
self.endpoint_queues : dict[str,Deque]
self.buffer_queues : dict[str,sortedlist]
For registering the implemented policies register at policy.py
Expected Behavior:
the Deque of all endpoints in self.endpoint_queues linked with
atleast one backend must be empty irrespective of whatever
backend policy is implemented.
"""
raise NotImplementedError(
"This method should be implemented by child class.")
# Flushes the specified backend queues and assigns work to workers.
async def flush_backend_queues(self, backends_to_flush):
assert self.flush_lock.locked()
for backend in backends_to_flush:
# No workers available.
if self.worker_queues[backend].qsize() == 0:
continue
# No work to do.
if len(self.backend_queues[backend]) == 0:
continue
# Flushes the buffer queue and assigns work to workers.
async def _flush_buffer_queues(self):
for endpoint in self.traffic:
ready_backends = self._get_available_backends(endpoint)
for backend in ready_backends:
# no work available
if len(self.buffer_queues[backend]) == 0:
continue
buffer_queue = self.backend_queues[backend]
worker_queue = self.worker_queues[backend]
buffer_queue = self.buffer_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), worker_queue.qsize()))
logger.debug("Assigning queries for backend {} with buffer "
"queue size {} and worker queue size {}".format(
backend, len(buffer_queue),
worker_queue.qsize()))
max_batch_size = None
if backend in self.backend_info:
max_batch_size = self.backend_info[backend].max_batch_size
max_batch_size = None
if backend in self.backend_info:
max_batch_size = self.backend_info[backend].max_batch_size
await self._assign_query_to_worker(
backend, buffer_queue, worker_queue, max_batch_size)
await self._assign_query_to_worker(backend, buffer_queue,
worker_queue, max_batch_size)
async def _do_query(self, backend, backend_replica_tag, req):
# If the worker died, this will be a RayActorError. Just return it and