def test_grpc_data_runtime(mocker):
args = set_pea_parser().parse_args([])
handle_mock = mocker.Mock()
def start_runtime(args, handle_mock):
with GRPCDataRuntime(args) as runtime:
runtime._data_request_handler.handle = handle_mock
runtime.run_forever()
runtime_thread = Thread(
target=start_runtime,
args=(
args,
handle_mock,
),
)
runtime_thread.start()
assert GRPCDataRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'{args.host}:{args.port_in}',
shutdown_event=Event())
Grpclet._create_grpc_stub(f'{args.host}:{args.port_in}',
is_async=False).Call(_create_test_data_message())
time.sleep(0.1)
handle_mock.assert_called()
GRPCDataRuntime.cancel(f'{args.host}:{args.port_in}')
runtime_thread.join()
assert not GRPCDataRuntime.is_ready(f'{args.host}:{args.port_in}')
def test_grpc_data_runtime(mocker):
args = set_pea_parser().parse_args([])
handle_mock = multiprocessing.Event()
cancel_event = multiprocessing.Event()
def start_runtime(args, handle_mock, cancel_event):
with GRPCDataRuntime(args, cancel_event) as runtime:
runtime._data_request_handler.handle = (
lambda *args, **kwargs: handle_mock.set()
)
runtime.run_forever()
runtime_thread = Process(
target=start_runtime,
args=(args, handle_mock, cancel_event),
daemon=True,
)
runtime_thread.start()
assert GRPCDataRuntime.wait_for_ready_or_shutdown(
timeout=5.0, ctrl_address=f'{args.host}:{args.port_in}', shutdown_event=Event()
)
Grpclet._create_grpc_stub(f'{args.host}:{args.port_in}', is_async=False).Call(
_create_test_data_message()
)
time.sleep(0.1)
assert handle_mock.is_set()
GRPCDataRuntime.cancel(cancel_event)
runtime_thread.join()
assert not GRPCDataRuntime.is_ready(f'{args.host}:{args.port_in}')
def test_grpc_data_runtime_waits_for_pending_messages_shutdown(close_method):
args = set_pea_parser().parse_args([])
cancel_event = multiprocessing.Event()
handler_closed_event = multiprocessing.Event()
slow_executor_block_time = 1.0
pending_requests = 3
sent_queue = multiprocessing.Queue()
def start_runtime(args, cancel_event, sent_queue, handler_closed_event):
with GRPCDataRuntime(args, cancel_event) as runtime:
runtime._data_request_handler.handle = lambda *args, **kwargs: time.sleep(
slow_executor_block_time
)
runtime._data_request_handler.close = (
lambda *args, **kwargs: handler_closed_event.set()
)
async def mock(msg):
sent_queue.put('')
runtime._grpclet.send_message = mock
runtime.run_forever()
runtime_thread = Process(
target=start_runtime,
args=(args, cancel_event, sent_queue, handler_closed_event),
daemon=True,
)
runtime_thread.start()
assert GRPCDataRuntime.wait_for_ready_or_shutdown(
timeout=5.0, ctrl_address=f'{args.host}:{args.port_in}', shutdown_event=Event()
)
request_start_time = time.time()
for i in range(pending_requests):
Grpclet._create_grpc_stub(f'{args.host}:{args.port_in}', is_async=False).Call(
_create_test_data_message()
)
time.sleep(0.1)
if close_method == 'TERMINATE':
runtime_thread.terminate()
else:
GRPCDataRuntime.cancel(cancel_event)
assert not handler_closed_event.is_set()
runtime_thread.join()
assert (
time.time() - request_start_time >= slow_executor_block_time * pending_requests
)
assert sent_queue.qsize() == pending_requests
assert handler_closed_event.is_set()
assert not GRPCDataRuntime.is_ready(f'{args.host}:{args.port_in}')
def test_executor_runtimes(signal, tmpdir, grpc_data_requests):
import time
args = set_pea_parser().parse_args([])
def run(args):
args.uses = {
'jtype': 'DummyExecutor',
'with': {
'dir': str(tmpdir)
},
'metas': {
'workspace': str(tmpdir)
},
}
args.grpc_data_requests = grpc_data_requests
executor_native(args)
process = multiprocessing.Process(target=run, args=(args, ))
process.start()
time.sleep(0.5)
if grpc_data_requests:
Grpclet._create_grpc_stub(f'{args.host}:{args.port_in}',
is_async=False).Call(
_create_test_data_message())
else:
socket = zmq.Context().socket(zmq.PUSH)
socket.connect(f'tcp://localhost:{args.port_in}')
socket.send_multipart(_create_test_data_message().dump())
time.sleep(0.1)
os.kill(process.pid, signal)
process.join()
with open(f'{tmpdir}/test.txt', 'r') as fp:
output = fp.read()
split = output.split(';')
assert split[0] == 'proper close'
assert split[1] == '1'