async def main_loop(self):
while True:
# NOTE(simon): There's an issue when user updated batch size and
# batch wait timeout during the execution, these values will not be
# updated until after the current iteration.
batch = await self.batch_queue.wait_for_batch()
all_evaluated_futures = []
if not self.config.accepts_batches:
query = batch[0]
evaluated = asyncio.ensure_future(self.invoke_single(query))
all_evaluated_futures = [evaluated]
chain_future(evaluated, query.async_future)
else:
get_call_method = attrgetter("call_method")
sorted_batch = sorted(batch, key=get_call_method)
for _, group in groupby(sorted_batch, key=get_call_method):
group = list(group)
evaluated = asyncio.ensure_future(self.invoke_batch(group))
all_evaluated_futures.append(evaluated)
result_futures = [q.async_future for q in group]
chain_future(unpack_future(evaluated, len(group)),
result_futures)
if self.config.is_blocking:
# We use asyncio.wait here so if the result is exception,
# it will not be raised.
await asyncio.wait(all_evaluated_futures)
async def main_loop(self) -> None:
while True:
# NOTE(simon): There's an issue when user updated batch size and
# batch wait timeout during the execution, these values will not be
# updated until after the current iteration.
batch = await self.batch_queue.wait_for_batch()
# Record metrics
self.num_queued_items.record(self.batch_queue.qsize(), {
"backend": self.backend_tag,
"replica_tag": self.replica_tag
})
self.num_processing_items.record(
self.num_ongoing_requests - self.batch_queue.qsize(), {
"backend": self.backend_tag,
"replica_tag": self.replica_tag
})
for query in batch:
queuing_time = (time.time() - query.tick_enter_replica) * 1000
self.queuing_latency_tracker.record(
queuing_time, {
"backend": self.backend_tag,
"replica_tag": self.replica_tag
})
all_evaluated_futures = []
if not self.config.internal_metadata.accepts_batches:
query = batch[0]
evaluated = asyncio.ensure_future(self.invoke_single(query))
all_evaluated_futures = [evaluated]
chain_future(evaluated, query.async_future)
else:
get_call_method = (
lambda query: query.metadata.call_method # noqa: E731
)
sorted_batch = sorted(batch, key=get_call_method)
for _, group in groupby(sorted_batch, key=get_call_method):
group = list(group)
evaluated = asyncio.ensure_future(self.invoke_batch(group))
all_evaluated_futures.append(evaluated)
result_futures = [q.async_future for q in group]
chain_future(unpack_future(evaluated, len(group)),
result_futures)
if self.config.internal_metadata.is_blocking:
# We use asyncio.wait here so if the result is exception,
# it will not be raised.
await asyncio.wait(all_evaluated_futures)
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 test_future_chaining():
def make():
return asyncio.get_event_loop().create_future()
# Test 1 -> 1 chaining
fut1, fut2 = make(), make()
chain_future(fut1, fut2)
fut1.set_result(1)
assert await fut2 == 1
# Test 1 -> 1 chaining with exception
fut1, fut2 = make(), make()
chain_future(fut1, fut2)
fut1.set_exception(ValueError(""))
with pytest.raises(ValueError):
await fut2
# Test many -> many chaining
src_futs = [make() for _ in range(4)]
dst_futs = [make() for _ in range(4)]
chain_future(src_futs, dst_futs)
[fut.set_result(i) for i, fut in enumerate(src_futs)]
for i, fut in enumerate(dst_futs):
assert await fut == i
# Test 1 -> many unwrapping
batched_future = make()
single_futures = unpack_future(batched_future, 4)
batched_future.set_result(list(range(4)))
for i, fut in enumerate(single_futures):
assert await fut == i
# Test 1 -> many unwrapping with exception
batched_future = make()
single_futures = unpack_future(batched_future, 4)
batched_future.set_exception(ValueError(""))
for future in single_futures:
with pytest.raises(ValueError):
await future
