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