def create_async_channel_stub(
address,
https=False,
root_certificates: Optional[str] = None,
) -> Tuple[jina_pb2_grpc.JinaSingleDataRequestRPCStub,
jina_pb2_grpc.JinaDataRequestRPCStub,
jina_pb2_grpc.JinaControlRequestRPCStub, grpc.aio.Channel, ]:
"""
Creates an async GRPC Channel. This channel has to be closed eventually!
:param address: the address to create the connection to, like 127.0.0.0.1:8080
:param https: if True, use https for the grpc channel
:param root_certificates: the path to the root certificates for https, only u
:returns: DataRequest/ControlRequest stubs and an async grpc channel
"""
channel = GrpcConnectionPool.get_grpc_channel(
address,
asyncio=True,
https=https,
root_certificates=root_certificates,
)
return (
jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel),
jina_pb2_grpc.JinaDataRequestRPCStub(channel),
jina_pb2_grpc.JinaControlRequestRPCStub(channel),
channel,
)
async def send_request_async(
request: Request,
target: str,
timeout: float = 1.0,
https: bool = False,
root_certificates: Optional[str] = None,
) -> Request:
"""
Sends a request asynchronously to the target via grpc
:param request: the request to send
:param target: where to send the request to, like 127.0.0.1:8080
:param timeout: timeout for the send
:param https: if True, use https for the grpc channel
:param root_certificates: the path to the root certificates for https, only u
:returns: the response request
"""
async with GrpcConnectionPool.get_grpc_channel(
target,
asyncio=True,
https=https,
root_certificates=root_certificates,
) as channel:
if type(request) == DataRequest:
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
return await stub.process_single_data(request, timeout=timeout)
elif type(request) == ControlRequest:
stub = jina_pb2_grpc.JinaControlRequestRPCStub(channel)
return await stub.process_control(request, timeout=timeout)
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 send_request_sync(
request: Request,
target: str,
timeout=100.0,
tls=False,
root_certificates: Optional[str] = None,
endpoint: Optional[str] = None,
) -> Request:
"""
Sends a request synchronously to the target via grpc
:param request: the request to send
:param target: where to send the request to, like 127.0.0.1:8080
:param timeout: timeout for the send
:param tls: if True, use tls encryption for the grpc channel
:param root_certificates: the path to the root certificates for tls, only used if tls is True
:param endpoint: endpoint to target with the request
:returns: the response request
"""
for i in range(3):
try:
with GrpcConnectionPool.get_grpc_channel(
target,
tls=tls,
root_certificates=root_certificates,
) as channel:
if type(request) == DataRequest:
metadata = (('endpoint',
endpoint), ) if endpoint else None
stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(
channel)
response, call = stub.process_single_data.with_call(
request,
timeout=timeout,
metadata=metadata,
)
elif type(request) == ControlRequest:
stub = jina_pb2_grpc.JinaControlRequestRPCStub(channel)
response = stub.process_control(request,
timeout=timeout)
return response
except grpc.RpcError as e:
if e.code() != grpc.StatusCode.UNAVAILABLE or i == 2:
raise
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}')
