def test_cycle():
graph = RoutingTable()
graph.add_pod('executor0', '0.0.0.0', 1230, 1232, '')
graph.add_pod('executor1', '0.0.0.0', 1231, 1233, '')
graph.add_edge('executor0', 'executor1')
graph.add_edge('executor1', 'executor0')
graph.active_pod = 'executor0'
assert not graph.is_acyclic()
def test_cycle():
graph = RoutingTable()
graph.add_pod('pod0', PodInterface('0.0.0.0', 1230, 1232))
graph.add_pod('pod1', PodInterface('0.0.0.0', 1231, 1233))
graph.add_edge('pod0', 'pod1')
graph.add_edge('pod1', 'pod0')
graph.active_pod = 'pod0'
assert not graph.is_acyclic()
def test_double_dynamic_routing_zmqlet():
args1 = get_args()
args2 = get_args()
args3 = get_args()
logger = JinaLogger('zmq-test')
with Zmqlet(args1, logger) as z1, Zmqlet(args2, logger) as z2, Zmqlet(
args3, logger
) as z3:
assert z1.msg_sent == 0
assert z2.msg_sent == 0
assert z3.msg_sent == 0
req = jina_pb2.RequestProto()
req.request_id = random_identity()
d = req.data.docs.add()
d.tags['id'] = 2
msg = Message(None, req, 'tmp', '')
routing_table = {
'active_pod': 'executor1',
'pods': {
'executor1': {
'host': __default_host__,
'port': args1.port_in,
'expected_parts': 0,
'out_edges': [{'pod': 'executor2'}, {'pod': 'executor3'}],
},
'executor2': {
'host': __default_host__,
'port': args2.port_in,
'expected_parts': 1,
'out_edges': [],
},
'executor3': {
'host': __default_host__,
'port': args3.port_in,
'expected_parts': 1,
'out_edges': [],
},
},
}
msg.envelope.routing_table.CopyFrom(RoutingTable(routing_table).proto)
number_messages = 100
trips = 10
for i in range(trips):
for j in range(number_messages):
z1.send_message(msg)
time.sleep(1)
for i in range(number_messages):
z2.recv_message(callback)
z3.recv_message(callback)
total_number_messages = number_messages * trips
assert z1.msg_sent == 2 * total_number_messages
assert z2.msg_sent == 0
assert z2.msg_recv == total_number_messages
assert z3.msg_sent == 0
assert z3.msg_recv == total_number_messages
def test_simple_routing():
graph = RoutingTable()
graph.add_pod('pod0', PodInterface('0.0.0.0', 1230, 1232))
graph.add_pod('pod1', PodInterface('0.0.0.0', 1231, 1233))
graph.add_edge('pod0', 'pod1')
graph.active_pod = 'pod0'
next_routes = graph.get_next_targets()
assert len(next_routes) == 1
assert next_routes[0][0].active_pod == 'pod1'
def test_simple_routing():
graph = RoutingTable()
graph.add_pod('executor0', '0.0.0.0', 1230, 1232, '')
graph.add_pod('executor1', '0.0.0.0', 1231, 1233, '')
graph.add_edge('executor0', 'executor1')
graph.active_pod = 'executor0'
next_routes = graph.get_next_targets()
assert len(next_routes) == 1
assert next_routes[0][0].active_pod == 'executor1'
async def send_message(self, msg: 'Message', **kwargs):
"""
Sends a message via gRPC to the target indicated in the message's routing table
:param msg: the protobuf message to send
:param kwargs: Additional arguments.
"""
if self._next_targets:
for pod_address in self._next_targets:
self._send_message(msg, pod_address)
else:
routing_table = RoutingTable(msg.envelope.routing_table)
next_targets = routing_table.get_next_targets()
for target, _ in next_targets:
self._send_message(
self._add_envelope(msg, target),
target.active_target_pod.full_address,
)
def _get_expected_parts(self, msg):
if msg.is_data_request:
if not self._static_routing_table:
graph = RoutingTable(msg.envelope.routing_table)
return graph.active_target_pod.expected_parts
else:
return self.args.num_part
else:
return 1
def __init__(
self,
args: argparse.Namespace,
message_callback: Callable[['Message'], None],
logger: Optional['JinaLogger'] = None,
):
self.args = args
self._stubs = {}
self._logger = logger or JinaLogger(self.__class__.__name__)
self.callback = message_callback
self.msg_recv = 0
self.msg_sent = 0
self._pending_tasks = []
self._send_routing_table = args.send_routing_table
if not args.send_routing_table:
self._routing_table = RoutingTable(args.routing_table)
self._next_targets = self._routing_table.get_next_target_addresses(
)
else:
self._routing_table = None
self._next_targets = None
def test_single_routing():
graph = RoutingTable()
graph.add_pod('pod0', PodInterface('0.0.0.0', 1230, 1233))
graph.active_pod = 'pod0'
next_routes = graph.get_next_targets()
assert len(next_routes) == 0
def test_single_routing():
graph = RoutingTable()
graph.add_pod('executor0', '0.0.0.0', 1230, 1233, '')
graph.active_pod = 'executor0'
next_routes = graph.get_next_targets()
assert len(next_routes) == 0
class Grpclet(jina_pb2_grpc.JinaDataRequestRPCServicer):
"""A `Grpclet` object can send/receive Messages via gRPC.
:param args: the parsed arguments from the CLI
:param message_callback: the callback to call on received messages
:param logger: the logger to use
"""
def __init__(
self,
args: argparse.Namespace,
message_callback: Callable[['Message'], None],
logger: Optional['JinaLogger'] = None,
):
self.args = args
self._logger = logger or JinaLogger(self.__class__.__name__)
self.callback = message_callback
self._connection_pool = create_connection_pool(args)
self.msg_recv = 0
self.msg_sent = 0
self._pending_tasks = []
self._static_routing_table = args.static_routing_table
if args.static_routing_table:
self._routing_table = RoutingTable(args.routing_table)
self._next_targets = self._routing_table.get_next_target_addresses()
else:
self._routing_table = None
self._next_targets = None
async def send_message(self, msg: 'Message', **kwargs):
"""
Sends a message via gRPC to the target indicated in the message's routing table
:param msg: the protobuf message to send
:param kwargs: Additional arguments.
"""
if self._next_targets:
for pod_address in self._next_targets:
self._send_message(msg, pod_address)
else:
routing_table = RoutingTable(msg.envelope.routing_table)
next_targets = routing_table.get_next_targets()
for target, _ in next_targets:
self._send_message(
self._add_envelope(msg, target),
target.active_target_pod.full_address,
)
def _send_message(self, msg, pod_address):
try:
self.msg_sent += 1
self._pending_tasks.append(
self._connection_pool.send_message(msg, pod_address)
)
self._update_pending_tasks()
except grpc.RpcError as ex:
self._logger.error('Sending data request via grpc failed', ex)
raise ex
@staticmethod
def send_ctrl_msg(pod_address: str, command: str, timeout=1.0):
"""
Sends a control message via gRPC to pod_address
:param pod_address: the pod to send the command to
:param command: the command to send (TERMINATE/ACTIVATE/...)
:param timeout: optional timeout for the request in seconds
:returns: Empty protobuf struct
"""
stub = Grpclet._create_grpc_stub(pod_address, is_async=False)
response = stub.Call(ControlMessage(command), timeout=timeout)
return response
@staticmethod
def _create_grpc_stub(pod_address, is_async=True):
if is_async:
channel = grpc.aio.insecure_channel(
pod_address,
options=[
('grpc.max_send_message_length', -1),
('grpc.max_receive_message_length', -1),
],
)
else:
channel = grpc.insecure_channel(
pod_address,
options=[
('grpc.max_send_message_length', -1),
('grpc.max_receive_message_length', -1),
],
)
stub = jina_pb2_grpc.JinaDataRequestRPCStub(channel)
return stub
def _add_envelope(self, msg, routing_table):
if not self._static_routing_table:
new_envelope = jina_pb2.EnvelopeProto()
new_envelope.CopyFrom(msg.envelope)
new_envelope.routing_table.CopyFrom(routing_table.proto)
new_message = Message(request=msg.request, envelope=new_envelope)
return new_message
else:
return msg
async def close(self, grace_period=None, *args, **kwargs):
"""Stop the Grpc server
:param grace_period: Time to wait for message processing to finish before killing the grpc server
:param args: Extra positional arguments
:param kwargs: Extra key-value arguments
"""
self._update_pending_tasks()
try:
await asyncio.wait_for(asyncio.gather(*self._pending_tasks), timeout=1.0)
except asyncio.TimeoutError:
self._update_pending_tasks()
self._logger.warning(
f'Could not gracefully complete {len(self._pending_tasks)} pending tasks on close.'
)
self._connection_pool.close()
self._logger.debug('Close grpc server')
await self._grpc_server.stop(grace_period)
async def start(self):
"""
Starts this Grpclet by starting its gRPC server
"""
self._grpc_server = grpc.aio.server(
options=[
('grpc.max_send_message_length', -1),
('grpc.max_receive_message_length', -1),
]
)
jina_pb2_grpc.add_JinaDataRequestRPCServicer_to_server(self, self._grpc_server)
bind_addr = f'{self.args.host}:{self.args.port_in}'
self._grpc_server.add_insecure_port(bind_addr)
self._logger.debug(f'Binding gRPC server for data requests to {bind_addr}')
self._connection_pool.start()
await self._grpc_server.start()
await self._grpc_server.wait_for_termination()
def _update_pending_tasks(self):
self._pending_tasks = [task for task in self._pending_tasks if not task.done()]
async def Call(self, msg, *args):
"""Processes messages received by the GRPC server
:param msg: The received message
:param args: Extra positional arguments
:return: Empty protobuf struct, necessary to return for protobuf Empty
"""
if self.callback:
self._pending_tasks.append(asyncio.create_task(self.callback(msg)))
else:
self._logger.debug(
'Grpclet received data request, but no callback was registered'
)
self.msg_recv += 1
self._update_pending_tasks()
return struct_pb2.Value()
def test_topological_sorting():
graph = RoutingTable()
graph.add_pod('executor0', '0.0.0.0', 1230, 1234, '')
graph.add_pod('executor1', '0.0.0.0', 1231, 1235, '')
graph.add_pod('executor2', '0.0.0.0', 1232, 1236, '')
graph.add_pod('executor3', '0.0.0.0', 1233, 1237, '')
graph.add_pod('executor4', '0.0.0.0', 1233, 1238, '')
graph.add_edge('executor0', 'executor1')
graph.add_edge('executor0', 'executor2')
graph.add_edge('executor1', 'executor3')
graph.add_edge('executor2', 'executor4')
graph.add_edge('executor3', 'executor4')
graph.active_pod = 'executor0'
topological_sorting = graph._topological_sort()
assert topological_sorting[0] == 'executor0'
assert topological_sorting[1] in ['executor1', 'executor2']
assert topological_sorting[2] in ['executor1', 'executor2', 'executor3']
assert topological_sorting[3] in ['executor2', 'executor3']
assert topological_sorting[4] == 'executor4'
def test_nested_routing():
graph = RoutingTable()
graph.add_pod('executor0', '0.0.0.0', 1230, 1234, '')
graph.add_pod('executor1', '0.0.0.0', 1231, 1235, '')
graph.add_pod('executor2', '0.0.0.0', 1232, 1236, '')
graph.add_pod('executor3', '0.0.0.0', 1233, 1237, '')
graph.add_pod('executor4', '0.0.0.0', 1233, 1238, '')
graph.add_edge('executor0', 'executor1')
graph.add_edge('executor0', 'executor2')
graph.add_edge('executor1', 'executor3')
graph.add_edge('executor2', 'executor4')
graph.add_edge('executor3', 'executor4')
graph.active_pod = 'executor0'
next_routes = graph.get_next_targets()
assert len(next_routes) == 2
assert next_routes[0][0].active_pod == 'executor1'
assert next_routes[1][0].active_pod == 'executor2'
graph.active_pod = 'executor1'
next_routes = graph.get_next_targets()
assert len(next_routes) == 1
assert next_routes[0][0].active_pod == 'executor3'
graph.active_pod = 'executor2'
next_routes = graph.get_next_targets()
assert len(next_routes) == 1
assert next_routes[0][0].active_pod == 'executor4'
graph.active_pod = 'executor3'
next_routes = graph.get_next_targets()
assert len(next_routes) == 1
assert next_routes[0][0].active_pod == 'executor4'
graph.active_pod = 'executor4'
next_routes = graph.get_next_targets()
assert len(next_routes) == 0
def test_no_cycle():
graph = RoutingTable()
graph.add_pod('pod0', PodInterface('0.0.0.0', 1230, 1234))
graph.add_pod('pod1', PodInterface('0.0.0.0', 1231, 1235))
graph.add_pod('pod2', PodInterface('0.0.0.0', 1232, 1236))
graph.add_pod('pod3', PodInterface('0.0.0.0', 1233, 1237))
graph.add_pod('pod4', PodInterface('0.0.0.0', 1233, 1238))
graph.add_edge('pod2', 'pod1')
graph.add_edge('pod1', 'pod0')
graph.add_edge('pod0', 'pod3')
graph.add_edge('pod3', 'pod4')
graph.active_pod = 'pod0'
assert graph.is_acyclic()
def test_no_cycle():
graph = RoutingTable()
graph.add_pod('executor0', '0.0.0.0', 1230, 1234, '')
graph.add_pod('executor1', '0.0.0.0', 1231, 1235, '')
graph.add_pod('executor2', '0.0.0.0', 1232, 1236, '')
graph.add_pod('executor3', '0.0.0.0', 1233, 1237, '')
graph.add_pod('executor4', '0.0.0.0', 1233, 1238, '')
graph.add_edge('executor2', 'executor1')
graph.add_edge('executor1', 'executor0')
graph.add_edge('executor0', 'executor3')
graph.add_edge('executor3', 'executor4')
graph.active_pod = 'executor0'
assert graph.is_acyclic()
def test_nested_routing():
graph = RoutingTable()
graph.add_pod('pod0', PodInterface('0.0.0.0', 1230, 1234))
graph.add_pod('pod1', PodInterface('0.0.0.0', 1231, 1235))
graph.add_pod('pod2', PodInterface('0.0.0.0', 1232, 1236))
graph.add_pod('pod3', PodInterface('0.0.0.0', 1233, 1237))
graph.add_pod('pod4', PodInterface('0.0.0.0', 1233, 1238))
graph.add_edge('pod0', 'pod1')
graph.add_edge('pod0', 'pod2')
graph.add_edge('pod1', 'pod3')
graph.add_edge('pod2', 'pod4')
graph.add_edge('pod3', 'pod4')
graph.active_pod = 'pod0'
next_routes = graph.get_next_targets()
assert len(next_routes) == 2
assert next_routes[0][0].active_pod == 'pod1'
assert next_routes[1][0].active_pod == 'pod2'
graph.active_pod = 'pod1'
next_routes = graph.get_next_targets()
assert len(next_routes) == 1
assert next_routes[0][0].active_pod == 'pod3'
graph.active_pod = 'pod2'
next_routes = graph.get_next_targets()
assert len(next_routes) == 1
assert next_routes[0][0].active_pod == 'pod4'
graph.active_pod = 'pod3'
next_routes = graph.get_next_targets()
assert len(next_routes) == 1
assert next_routes[0][0].active_pod == 'pod4'
graph.active_pod = 'pod4'
next_routes = graph.get_next_targets()
assert len(next_routes) == 0
def test_topological_sorting():
graph = RoutingTable()
graph.add_pod('pod0', PodInterface('0.0.0.0', 1230, 1234))
graph.add_pod('pod1', PodInterface('0.0.0.0', 1231, 1235))
graph.add_pod('pod2', PodInterface('0.0.0.0', 1232, 1236))
graph.add_pod('pod3', PodInterface('0.0.0.0', 1233, 1237))
graph.add_pod('pod4', PodInterface('0.0.0.0', 1233, 1238))
graph.add_edge('pod0', 'pod1')
graph.add_edge('pod0', 'pod2')
graph.add_edge('pod1', 'pod3')
graph.add_edge('pod2', 'pod4')
graph.add_edge('pod3', 'pod4')
graph.active_pod = 'pod0'
topological_sorting = graph._topological_sort()
assert topological_sorting[0] == 'pod0'
assert topological_sorting[1] in ['pod1', 'pod2']
assert topological_sorting[2] in ['pod1', 'pod2', 'pod3']
assert topological_sorting[3] in ['pod2', 'pod3']
assert topological_sorting[4] == 'pod4'
def ip_from(flow, pod_number):
return RoutingTable(
flow['gateway'].args.routing_table).pods[pod_number].host