def test_simple_zmqlet():
args = set_pea_parser().parse_args([
'--host-in', '0.0.0.0',
'--host-out', '0.0.0.0',
'--port-in', '12346',
'--port-out', '12347',
'--socket-in', 'PULL_CONNECT',
'--socket-out', 'PUSH_CONNECT',
'--timeout-ctrl', '-1'])
args2 = set_pea_parser().parse_args([
'--host-in', '0.0.0.0',
'--host-out', '0.0.0.0',
'--port-in', '12347',
'--port-out', '12346',
'--socket-in', 'PULL_BIND',
'--socket-out', 'PUSH_BIND',
'--uses', '_logforward',
'--timeout-ctrl', '-1'
])
logger = logging.getLogger('zmq-test')
with BasePea(args2) as z1, Zmqlet(args, logger) as z:
req = jina_pb2.RequestProto()
req.request_id = get_random_identity()
d = req.index.docs.add()
d.tags['id'] = 2
msg = Message(None, req, 'tmp', '')
z.send_message(msg)
def test_double_dynamic_routing_zmqstreamlet():
args1 = get_args()
args2 = get_args()
args3 = get_args()
logger = JinaLogger('zmq-test')
with ZmqStreamlet(args=args1, logger=logger) as z1, ZmqStreamlet(
args=args2, logger=logger
) as z2, ZmqStreamlet(args=args3, logger=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_pb = jina_pb2.RoutingTableProto()
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': [],
},
},
}
json_format.ParseDict(routing_table, routing_pb)
msg.envelope.routing_table.CopyFrom(routing_pb)
for pea in [z1, z2, z3]:
thread = threading.Thread(target=pea.start, args=(callback,))
thread.daemon = True
thread.start()
number_messages = 1000
for i in range(number_messages):
z1.send_message(msg)
time.sleep(5)
assert z1.msg_sent == 2 * number_messages
assert z1.msg_recv == 0
assert z2.msg_sent == 0
assert z2.msg_recv == number_messages
assert z3.msg_sent == 0
assert z3.msg_recv == number_messages
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_not_read_zmqlet():
with MockBasePeaNotRead(args3), Zmqlet(args1, default_logger) as z:
req = jina_pb2.RequestProto()
req.request_id = get_random_identity()
d = req.index.docs.add()
d.tags['id'] = 2
msg = Message(None, req, 'tmp', '')
z.send_message(msg)
def test_read_zmqlet():
with MockBasePeaRead(args2), Zmqlet(args1, default_logger) as z:
req = jina_pb2.RequestProto()
req.request_id = uuid.uuid1().hex
d = req.index.docs.add()
d.tags['id'] = 2
msg = Message(None, req, 'tmp', '')
z.send_message(msg)
def test_add_route():
r = jina_pb2.RequestProto()
r.control.command = jina_pb2.RequestProto.ControlRequestProto.IDLE
msg = Message(None, r, pod_name='test1', identity='sda')
msg.add_route('name', 'identity')
assert len(msg.envelope.routes) == 2
assert msg.envelope.routes[1].pod == 'name'
assert msg.envelope.routes[1].pod_id == 'identity'
async def test_double_dynamic_routing_async_zmqlet():
args1 = get_args()
args2 = get_args()
args3 = get_args()
logger = JinaLogger('zmq-test')
with AsyncZmqlet(args1, logger) as z1, AsyncZmqlet(
args2, logger) as z2, AsyncZmqlet(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_pb = jina_pb2.RoutingTableProto()
routing_table = {
'active_pod': 'pod1',
'pods': {
'pod1': {
'host': '0.0.0.0',
'port': args1.port_in,
'expected_parts': 0,
'out_edges': [{
'pod': 'pod2'
}, {
'pod': 'pod3'
}],
},
'pod2': {
'host': '0.0.0.0',
'port': args2.port_in,
'expected_parts': 1,
'out_edges': [],
},
'pod3': {
'host': '0.0.0.0',
'port': args3.port_in,
'expected_parts': 1,
'out_edges': [],
},
},
}
json_format.ParseDict(routing_table, routing_pb)
msg.envelope.routing_table.CopyFrom(routing_pb)
await send_msg(z1, msg)
await z2.recv_message(callback)
await z3.recv_message(callback)
assert z1.msg_sent == 2
assert z1.msg_recv == 0
assert z2.msg_sent == 0
assert z2.msg_recv == 1
assert z3.msg_sent == 0
assert z3.msg_recv == 1
def eval_request():
num_docs = 10
req = jina_pb2.RequestProto()
for idx in range(num_docs):
doc = Document(req.index.docs.add())
gt = Document(req.index.groundtruths.add())
chunk_doc = doc.chunks.new()
chunk_gt = gt.chunks.new()
chunk_doc.embedding = np.array([1, 1])
chunk_gt.embedding = np.array([2, 2])
return Request(req)
def test_remote_local_dynamic_routing_zmqlet():
args1 = get_args()
args2 = get_args()
args2.zmq_identity = 'test-identity'
args2.hosts_in_connect = [f'{args1.host}:{args1.port_out}']
logger = JinaLogger('zmq-test')
with Zmqlet(args1, logger) as z1, Zmqlet(args2, logger) as z2:
assert z1.msg_sent == 0
assert z1.msg_recv == 0
assert z2.msg_sent == 0
assert z2.msg_recv == 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_pb = jina_pb2.RoutingTableProto()
routing_table = {
'active_pod': 'pod1',
'pods': {
'pod1': {
'host': '0.0.0.0',
'port': args1.port_in,
'expected_parts': 0,
'out_edges': [{
'pod': 'pod2',
'send_as_bind': True
}],
},
'pod2': {
'host': '0.0.0.0',
'port': args2.port_in,
'expected_parts': 1,
'out_edges': [],
'target_identity': args2.zmq_identity,
},
},
}
json_format.ParseDict(routing_table, routing_pb)
msg.envelope.routing_table.CopyFrom(routing_pb)
z2.recv_message(callback)
assert z2.msg_sent == 0
assert z2.msg_recv == 0
z1.send_message(msg)
z2.recv_message(callback)
assert z1.msg_sent == 1
assert z1.msg_recv == 0
assert z2.msg_sent == 0
assert z2.msg_recv == 1
def test_simple_dynamic_routing_zmqlet():
args1 = get_args()
args2 = get_args()
logger = JinaLogger('zmq-test')
with Zmqlet(args1, logger) as z1, Zmqlet(args2, logger) as z2:
assert z1.msg_sent == 0
assert z1.msg_recv == 0
assert z2.msg_sent == 0
assert z2.msg_recv == 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_pb = jina_pb2.RoutingTableProto()
routing_table = {
'active_pod': 'executor1',
'pods': {
'executor1': {
'host': __default_host__,
'port': args1.port_in,
'expected_parts': 0,
'out_edges': [{'pod': 'executor2'}],
},
'executor2': {
'host': __default_host__,
'port': args2.port_in,
'expected_parts': 1,
'out_edges': [],
},
},
}
json_format.ParseDict(routing_table, routing_pb)
msg.envelope.routing_table.CopyFrom(routing_pb)
z2.recv_message(callback)
assert z2.msg_sent == 0
assert z2.msg_recv == 0
z1.send_message(msg)
z2.recv_message(callback)
assert z1.msg_sent == 1
assert z1.msg_recv == 0
assert z2.msg_sent == 0
assert z2.msg_recv == 1
def test_simple_zmqlet():
args = set_pea_parser().parse_args(
[
'--host-in',
'0.0.0.0',
'--host-out',
'0.0.0.0',
'--socket-in',
'PULL_CONNECT',
'--socket-out',
'PUSH_CONNECT',
'--timeout-ctrl',
'-1',
]
)
args2 = set_pea_parser().parse_args(
[
'--host-in',
'0.0.0.0',
'--host-out',
'0.0.0.0',
'--port-in',
str(args.port_out),
'--port-out',
str(args.port_in),
'--socket-in',
'PULL_BIND',
'--socket-out',
'PUSH_BIND',
'--timeout-ctrl',
'-1',
]
)
logger = JinaLogger('zmq-test')
with BasePea(args2), Zmqlet(args, logger) as z:
req = jina_pb2.RequestProto()
req.request_id = random_identity()
d = req.data.docs.add()
d.tags['id'] = 2
msg = Message(None, req, 'tmp', '')
z.send_message(msg)
def request(field_type):
num_docs = 10
req = jina_pb2.RequestProto()
for idx in range(num_docs):
doc = req.index.docs.add()
gt = req.index.groundtruths.add()
chunk_doc = Document(doc.chunks.add())
chunk_gt = Document(gt.chunks.add())
chunk_doc.granularity = 1
chunk_gt.granularity = 1
if field_type == 'text':
chunk_doc.text = 'aaa'
chunk_gt.text = 'aaaa'
elif field_type == 'buffer':
chunk_doc.buffer = b'\x01\x02\x03'
chunk_gt.buffer = b'\x01\x02\x03\x04'
elif field_type == 'blob':
chunk_doc.blob = np.array([1, 1, 1])
chunk_gt.blob = np.array([1, 1, 1, 1])
return Request(req).as_typed_request('index')
def req():
r = jina_pb2.RequestProto()
r.request_id = get_random_identity()
r.index.docs.add()
return r
def _generate_request():
req = jina_pb2.RequestProto()
req.request_id = random_identity()
req.data.docs.add()
return req
def req():
r = jina_pb2.RequestProto()
r.request_id = random_identity()
r.data.docs.add()
return r
