async def _flush_buffer_queues(self):
for service in self.traffic.keys():
ready_backends = self._get_available_backends(service)
for backend in ready_backends:
# no work available
if len(self.buffer_queues[backend]) == 0:
continue
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()
)
)
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
)
async def _flush_service_queues(self):
# perform traffic splitting for requests
for service, queue in self.service_queues.items():
# while there are incoming requests and there are backends
while queue.qsize() and len(self.traffic[service]):
backend_names = list(self.traffic[service].keys())
backend_weights = list(self.traffic[service].values())
if len(self.traffic[service]) >= 2:
# randomly pick 2 backends
backend1, backend2 = np.random.choice(backend_names,
2,
replace=False,
p=backend_weights)
# see the length of buffer queues of the two backends
# and pick the one which has less no. of queries
# in the buffer
if len(self.buffer_queues[backend1]) <= len(
self.buffer_queues[backend2]):
chosen_backend = backend1
else:
chosen_backend = backend2
logger.debug("[Power of two chocies] found two backends "
"{} and {}: choosing {}.".format(
backend1, backend2, chosen_backend))
else:
chosen_backend = np.random.choice(
backend_names, replace=False,
p=backend_weights).squeeze()
request = await queue.get()
self.buffer_queues[chosen_backend].add(request)
async def set_traffic(self, service, traffic_dict):
logger.debug("Setting traffic for service %s to %s", service,
traffic_dict)
self.traffic[service] = traffic_dict
backend_names = list(self.traffic[service].keys())
self.round_robin_iterator_map[service] = itertools.cycle(backend_names)
await self.flush()
async def set_traffic(self, service, traffic_dict):
logger.debug("Setting traffic for service %s to %s", service,
traffic_dict)
self.traffic[service] = traffic_dict
backend_names = list(self.traffic[service].keys())
self.fixed_packing_iterator_map[service] = itertools.cycle(
itertools.chain.from_iterable(
itertools.repeat(x, self.packing_num) for x in backend_names))
await self.flush()
async def enqueue_request(
self, request_meta, *request_args, **request_kwargs
):
service = request_meta.service
logger.debug("Received a request for service {}".format(service))
query = self._make_query(request_meta, request_args, request_kwargs)
query.on_enqueue(service, metadata=request_meta.tracing_metadata)
await self.service_queues[service].put(query)
await self.flush()
# Note: a future change can be to directly return the ObjectID from
# replica task submission
result = await query.async_future
# asyncio.get_event_loop().create_task(query.on_complete(self))
return result
def register_service(self, route: Union[str, None], service: str,
methods: List[str]):
"""Create an entry in the routing table
Args:
route: http path name. Must begin with '/'.
service: service name. This is the name http actor will push
the request to.
"""
logger.debug(
"[KV] Registering route {} to service {} with methods {}.".format(
route, service, methods))
# put no route services in default key
if route is None:
no_http_services = json.loads(
self.routing_table.get(NO_ROUTE_KEY, "[]"))
no_http_services.append(service)
self.routing_table.put(NO_ROUTE_KEY, json.dumps(no_http_services))
else:
self.routing_table.put(route, service)
self.methods_table.put(route, json.dumps(methods))
async def set_backend_config(self, backend, config_dict):
logger.debug(
"Setting backend config for "
"backend {} to {}".format(backend, config_dict)
)
self.backend_info[backend] = config_dict
async def set_traffic(self, service, traffic_dict):
logger.debug(
"Setting traffic for service %s to %s", service, traffic_dict
)
self.traffic[service] = traffic_dict
await self.flush()
async def link(self, service, backend):
logger.debug("Link %s with %s", service, backend)
await self.set_traffic(service, {backend: 1.0})
async def dequeue_request(self, backend, replica_handle):
logger.debug(
"Received a dequeue request for backend {}".format(backend)
)
await self.worker_queues[backend].put(replica_handle)
await self.flush()