def test_dynamic_polling(polling):
args = set_pod_parser().parse_args([
'--polling',
json.dumps({
'/any': PollingType.ANY,
'/all': PollingType.ALL,
'*': polling
}),
'--shards',
str(2),
])
connection_list_dict = {0: [f'fake_ip:8080'], 1: [f'fake_ip:8080']}
args.connection_list = json.dumps(connection_list_dict)
cancel_event, handle_queue, runtime_thread = _create_runtime(args)
with grpc.insecure_channel(
f'{args.host}:{args.port}',
options=GrpcConnectionPool.get_default_grpc_options(),
) as channel:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
response, call = stub.process_single_data.with_call(
_create_test_data_message(endpoint='all'),
metadata=(('endpoint', '/all'), ))
assert response
assert _queue_length(handle_queue) == 2
with grpc.insecure_channel(
f'{args.host}:{args.port}',
options=GrpcConnectionPool.get_default_grpc_options(),
) as channel:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
response, call = stub.process_single_data.with_call(
_create_test_data_message(endpoint='any'),
metadata=(('endpoint', '/any'), ))
assert response
assert _queue_length(handle_queue) == 3
_destroy_runtime(args, cancel_event, runtime_thread)
async def test_worker_runtime_slow_async_exec(uses):
args = set_pod_parser().parse_args(['--uses', uses])
cancel_event = multiprocessing.Event()
def start_runtime(args, cancel_event):
with WorkerRuntime(args, cancel_event=cancel_event) as runtime:
runtime.run_forever()
runtime_thread = Process(
target=start_runtime,
args=(args, cancel_event),
daemon=True,
)
runtime_thread.start()
assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'{args.host}:{args.port}',
ready_or_shutdown_event=Event(),
)
target = f'{args.host}:{args.port}'
results = []
async with grpc.aio.insecure_channel(
target,
options=GrpcConnectionPool.get_default_grpc_options(),
) as channel:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
tasks = []
for i in range(10):
async def task_wrapper():
return await stub.process_single_data(
_create_test_data_message())
tasks.append(asyncio.create_task(task_wrapper()))
for future in asyncio.as_completed(tasks):
t = await future
results.append(t.docs[0].text)
cancel_event.set()
runtime_thread.join()
if uses == 'AsyncSlowNewDocsExecutor':
assert results == ['1', '3', '5', '7', '9', '2', '4', '6', '8', '10']
else:
assert results == ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10']
assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')
def test_base_polling(polling):
args = set_pod_parser().parse_args([
'--polling',
polling,
'--shards',
str(2),
])
cancel_event, handle_queue, runtime_thread = _create_runtime(args)
_add_worker(args, shard_id=0)
_add_worker(args, shard_id=1)
with grpc.insecure_channel(
f'{args.host}:{args.port}',
options=GrpcConnectionPool.get_default_grpc_options(),
) as channel:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
response, call = stub.process_single_data.with_call(
_create_test_data_message(endpoint='all'),
metadata=(('endpoint', '/all'), ))
assert response
assert _queue_length(handle_queue) == 2 if polling == 'all' else 1
with grpc.insecure_channel(
f'{args.host}:{args.port}',
options=GrpcConnectionPool.get_default_grpc_options(),
) as channel:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
response, call = stub.process_single_data.with_call(
_create_test_data_message(endpoint='any'),
metadata=(('endpoint', '/any'), ))
assert response
assert _queue_length(handle_queue) == 4 if polling == 'all' else 2
_destroy_runtime(args, cancel_event, runtime_thread)
def test_error_in_worker_runtime(monkeypatch):
args = set_pod_parser().parse_args([])
cancel_event = multiprocessing.Event()
def fail(*args, **kwargs):
raise RuntimeError('intentional error')
monkeypatch.setattr(DataRequestHandler, 'handle', fail)
def start_runtime(args, cancel_event):
with WorkerRuntime(args, cancel_event=cancel_event) as runtime:
runtime.run_forever()
runtime_thread = Process(
target=start_runtime,
args=(args, cancel_event),
daemon=True,
)
runtime_thread.start()
assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'{args.host}:{args.port}',
ready_or_shutdown_event=Event(),
)
target = f'{args.host}:{args.port}'
with grpc.insecure_channel(
target,
options=GrpcConnectionPool.get_default_grpc_options(),
) as channel:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
response, call = stub.process_single_data.with_call(
_create_test_data_message())
assert response.header.status.code == jina_pb2.StatusProto.ERROR
assert 'is-error' in dict(call.trailing_metadata())
cancel_event.set()
runtime_thread.join()
assert response
assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')
def test_timeout_behaviour():
args = set_pod_parser().parse_args(['--timeout-send', '100'])
args.polling = PollingType.ANY
cancel_event, handle_queue, runtime_thread = _create_runtime(args)
_add_worker(args)
with grpc.insecure_channel(
f'{args.host}:{args.port}',
options=GrpcConnectionPool.get_default_grpc_options(),
) as channel:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
response, call = stub.process_single_data.with_call(
_create_test_data_message())
assert response
assert 'is-error' in dict(call.trailing_metadata())
assert len(response.docs) == 1
assert not handle_queue.empty()
_destroy_runtime(args, cancel_event, runtime_thread)
def test_regular_data_case():
args = set_pod_parser().parse_args([])
args.polling = PollingType.ANY
connection_list_dict = {0: [f'fake_ip:8080']}
args.connection_list = json.dumps(connection_list_dict)
cancel_event, handle_queue, runtime_thread = _create_runtime(args)
with grpc.insecure_channel(
f'{args.host}:{args.port}',
options=GrpcConnectionPool.get_default_grpc_options(),
) as channel:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
response, call = stub.process_single_data.with_call(
_create_test_data_message())
assert response
assert 'is-error' in dict(call.trailing_metadata())
assert len(response.docs) == 1
assert not handle_queue.empty()
_destroy_runtime(args, cancel_event, runtime_thread)
def test_decompress(monkeypatch):
call_counts = multiprocessing.Manager().Queue()
def decompress(self):
call_counts.put_nowait('called')
from jina.proto import jina_pb2
self._pb_body = jina_pb2.DataRequestProto()
self._pb_body.ParseFromString(self.buffer)
self.buffer = None
monkeypatch.setattr(
DataRequest,
'_decompress',
decompress,
)
args = set_pod_parser().parse_args([])
args.polling = PollingType.ANY
cancel_event, handle_queue, runtime_thread = _create_runtime(args)
_add_worker(args)
with grpc.insecure_channel(
f'{args.host}:{args.port}',
options=GrpcConnectionPool.get_default_grpc_options(),
) as channel:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
response, call = stub.process_single_data.with_call(
_create_test_data_message())
assert response
assert 'is-error' in dict(call.trailing_metadata())
assert _queue_length_copy(call_counts) == 0
assert len(response.docs) == 1
assert _queue_length_copy(call_counts) == 1
assert not handle_queue.empty()
_destroy_runtime(args, cancel_event, runtime_thread)
def test_worker_runtime():
args = set_pod_parser().parse_args([])
cancel_event = multiprocessing.Event()
def start_runtime(args, cancel_event):
with WorkerRuntime(args, cancel_event=cancel_event) as runtime:
runtime.run_forever()
runtime_thread = Process(
target=start_runtime,
args=(args, cancel_event),
daemon=True,
)
runtime_thread.start()
assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'{args.host}:{args.port}',
ready_or_shutdown_event=Event(),
)
target = f'{args.host}:{args.port}'
with grpc.insecure_channel(
target,
options=GrpcConnectionPool.get_default_grpc_options(),
) as channel:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
response, call = stub.process_single_data.with_call(
_create_test_data_message())
cancel_event.set()
runtime_thread.join()
assert response
assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')