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}')