def test_address_in_use(runtime): p = ['--port-ctrl', '55555', '--runtime-backend', runtime] args1 = set_pea_parser().parse_args(p) args2 = set_pea_parser().parse_args(p) with pytest.raises(RuntimeFailToStart): with Pea(args1), Pea(args2): pass
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_address_in_use(runtime): # test does not work with same control port, because right now the `READYNESS` signal is done by sending to that port p = ['--port-in', '55555', '--runtime-backend', runtime] args1 = set_pea_parser().parse_args(p) args2 = set_pea_parser().parse_args(p) with pytest.raises(RuntimeFailToStart): args1.name = 'Pea-1' args2.name = 'Pea-2' with Pea(args1), Pea(args2): pass
async def test_grpc_connection_pool(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, ), daemon=True, ) runtime_thread.start() assert GRPCDataRuntime.wait_for_ready_or_shutdown( timeout=3.0, ctrl_address=f'{args.host}:{args.port_in}', shutdown_event=Event() ) pool = GrpcConnectionPool() await pool.send_message( msg=_create_test_data_message(), target_address=f'{args.host}:{args.port_in}' ) time.sleep(0.1) handle_mock.assert_called() pool.close() GRPCDataRuntime.cancel(f'{args.host}:{args.port_in}') runtime_thread.join()
def test_base_pea_with_runtime_bad_teardown(mocker): class Pea1(Pea): def __init__(self, args): super().__init__(args) def mock_run_forever(*args, **kwargs): time.sleep(3) def mock_is_ready(*args, **kwargs): return True def mock_cancel(*args, **kwargs): pass mocker.patch.object(ZEDRuntime, 'run_forever', mock_run_forever) mocker.patch.object(ZEDRuntime, 'is_ready', mock_is_ready) mocker.patch.object(ZEDRuntime, 'teardown', lambda x: bad_func) mocker.patch.object(ZEDRuntime, 'cancel', lambda *args, **kwargs: mock_cancel) teardown_spy = mocker.spy(ZEDRuntime, 'teardown') cancel_spy = mocker.spy(Pea, '_cancel_runtime') run_spy = mocker.spy(ZEDRuntime, 'run_forever') arg = set_pea_parser().parse_args(['--runtime-backend', 'thread']) with Pea1(arg): pass teardown_spy.assert_called() run_spy.assert_called() cancel_spy.assert_called_once() # 3s > .join(1), need to cancel
def test_data_request_handler_change_docs_dam(logger, tmpdir): class MemmapExecutor(Executor): @requests def foo(self, docs, **kwargs): dam = DocumentArrayMemmap(tmpdir + '/dam') dam.extend(docs) return dam args = set_pea_parser().parse_args(['--uses', 'MemmapExecutor']) handler = DataRequestHandler(args, logger) req = list( request_generator( '/', DocumentArray([Document(text='input document') for _ in range(10)])))[0] msg = Message(None, req, 'test', '123') assert len(msg.request.docs) == 10 handler.handle( msg=msg, partial_requests=None, peapod_name='name', ) assert len(msg.request.docs) == 10 for doc in msg.request.docs: assert doc.text == 'input document'
def test_load_pod_with_custom_driver(): args = set_pea_parser().parse_args( ['--uses', str(cur_dir / 'yaml/test-executor-with-custom-driver.yml')]) with BasePea(args): # load success with no error pass
def test_base_pea_with_runtime_bad_cancel(mocker): class Pea1(Pea): def __init__(self, args): super().__init__(args) def mock_run_forever(runtime): time.sleep(3) def mock_is_ready(*args, **kwargs): return True mocker.patch.object(ZEDRuntime, 'run_forever', mock_run_forever) mocker.patch.object(ZEDRuntime, 'is_ready', mock_is_ready) mocker.patch.object(Pea, '_cancel_runtime', bad_func) teardown_spy = mocker.spy(ZEDRuntime, 'teardown') cancel_spy = mocker.spy(Pea, '_cancel_runtime') run_spy = mocker.spy(ZEDRuntime, 'run_forever') arg = set_pea_parser().parse_args(['--runtime-backend', 'thread']) with Pea1(arg): time.sleep(0.1) pass teardown_spy.assert_called() run_spy.assert_called() cancel_spy.assert_called_once()
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 _pea_command(self) -> str: from jina.parsers import set_pea_parser non_defaults = ArgNamespace.get_non_defaults_args(self.args, set_pea_parser(), taboo={'host'}) _args = ArgNamespace.kwargs2list(non_defaults) return f'jina pea {" ".join(_args)}'
def test_pea_runtime_env_setting_in_thread(fake_env): class EnvChecker(BaseExecutor): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # pea/pod-specific assert 'key1' not in os.environ assert 'key2' not in os.environ # inherit from parent process assert os.environ['key_parent'] == 'value3' os.environ['key_parent'] = 'value3' with Pea(set_pea_parser().parse_args([ '--uses', 'EnvChecker', '--env', 'key1=value1', '--env', 'key2=value2', '--runtime-backend', 'thread', ])): pass # should not affect the main process assert 'key1' not in os.environ assert 'key2' not in os.environ assert 'key_parent' in os.environ os.unsetenv('key_parent')
async def test_send_static_ctrl_msg(mocker): # AsyncMock does not seem to exist in python 3.7, this is a manual workaround receive_cb = mocker.Mock() async def mock_wrapper(msg): receive_cb() args = set_pea_parser().parse_args([]) grpclet = Grpclet(args=args, message_callback=mock_wrapper) asyncio.get_event_loop().create_task(grpclet.start()) receive_cb.assert_not_called() while True: try: def send_status(): return Grpclet.send_ctrl_msg( pod_address=f'{args.host}:{args.port_in}', command='STATUS') await asyncio.get_event_loop().run_in_executor(None, send_status) break except RpcError: await asyncio.sleep(0.1) receive_cb.assert_called() await grpclet.close(None)
async def test_grpc_connection_pool(): args = set_pea_parser().parse_args([]) is_called = multiprocessing.Event() is_cancel = Event() def start_runtime(args, is_called, is_cancel): with GRPCDataRuntime(args, is_cancel) as runtime: runtime._data_request_handler.handle = ( lambda *args, **kwargs: is_called.set() ) runtime.run_forever() runtime_process = Process( target=start_runtime, args=(args, is_called, is_cancel), ) runtime_process.start() assert GRPCDataRuntime.wait_for_ready_or_shutdown( timeout=3.0, ctrl_address=f'{args.host}:{args.port_in}', shutdown_event=Event() ) pool = GrpcConnectionPool() await pool.send_message( msg=_create_test_data_message(), target_address=f'{args.host}:{args.port_in}' ) time.sleep(0.1) assert is_called.is_set() pool.close() GRPCDataRuntime.cancel(cancel_event=is_cancel) runtime_process.terminate() runtime_process.join()
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_ssh_pea(): p = set_pea_parser().parse_args( ['--host', '[email protected]', '--timeout', '5000']) with SSHRuntime(p, kind='pea') as pp: assert pp.status.envelope.status.code == jina_pb2.StatusProto.READY assert pp.status is None
def test_simple_container_with_ext_yaml(docker_image_built): args = set_pea_parser().parse_args([ '--uses', f'docker://{img_name}', '--uses-internal', os.path.join(cur_dir, '../../../mwu-encoder/mwu_encoder_ext.yml') ]) with Pea(args): time.sleep(2)
def test_class_yaml3(): args = set_pea_parser().parse_args([]) with BasePea(args): pass from jina.executors.requests import _defaults assert _defaults is not None
def test_simple_container(docker_image_built): args = set_pea_parser().parse_args(['--uses', f'docker://{img_name}']) with Pea(args): pass time.sleep(2) Pea(args).start().close()
def test_pea_instantiate_start_same_context(): arg = set_pea_parser().parse_args([]) peas_args = [arg, arg] for args in peas_args: pea = Pea(args) with pea: pass
def test_peapod_create(mocker, api): args = set_pea_parser().parse_args([]) mocker.return_value.status_code = 201 mocker.return_value.json.return_value = 'abcd' assert api.create(args) == 'abcd' mocker.return_value.status_code = 404 assert not api.create(args)
def test_pea_store(): args = set_pea_parser().parse_args([]) store = InMemoryPeaStore() with store._session(): pea_id = store._create(pea_arguments=args) assert pea_id in store._store.keys() # assert isinstance(store._store[pea_id]['pea'], LocalRuntime) store._delete(pea_id) assert pea_id not in store._store.keys()
def test_pass_arbitrary_kwargs(monkeypatch, mocker): import docker mock = mocker.Mock() class MockContainers: def __init__(self): pass def run(self, *args, **kwargs): mock_kwargs = { k: kwargs[k] for k in ['hello', 'ports', 'environment'] } mock(**mock_kwargs) assert 'ports' in kwargs assert kwargs['ports'] is None assert 'environment' in kwargs assert kwargs['environment'] == ['VAR1=BAR', 'VAR2=FOO'] assert 'hello' in kwargs assert kwargs['hello'] == 0 class MockClient: def __init__(self, *args, **kwargs): pass def close(self): pass @property def networks(self): return {'bridge': None} @property def containers(self): return MockContainers() @property def images(self): return {} monkeypatch.setattr(docker, 'from_env', MockClient) args = set_pea_parser().parse_args([ '--uses', 'docker://jinahub/pod', '--docker-kwargs', 'hello: 0', 'environment: ["VAR1=BAR", "VAR2=FOO"]', ]) runtime = ContainerRuntime(args) runtime._docker_run(replay=False) expected_args = { 'hello': 0, 'ports': None, 'environment': ['VAR1=BAR', 'VAR2=FOO'] } mock.assert_called_with(**expected_args)
def test_peastore_add(partial_pea_store): partial_store_item = partial_pea_store.add( args=ArgNamespace.kwargs2namespace(PeaModel().dict(), set_pea_parser())) assert partial_store_item assert partial_pea_store.object assert partial_store_item.arguments['runtime_cls'] == 'ZEDRuntime' assert partial_store_item.arguments['host_in'] == __default_host__ assert partial_store_item.arguments['host_out'] == __default_host__
def test_container_runtime_good_entrypoint(runtime): class Pea1(BasePea): runtime_cls = ContainerRuntime arg = set_pea_parser().parse_args(['--uses', 'docker://jinaai/jina:test-pip', '--entrypoint', 'jina pod', '--runtime-backend', runtime]) with Pea1(arg): pass
def test_pea_set_shard_pea_id(): args = set_pea_parser().parse_args(['--shard-id', '1', '--shards', '3']) pea = Pea(args) assert pea.args.shard_id == 1 assert pea.args.pea_id == 1 assert pea.args.shards == 3 assert pea.args.parallel == 3
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_pass_native_arg(monkeypatch, mocker): import docker mocker.patch( 'jina.peapods.runtimes.container.ContainerRuntime.is_ready', return_value=True, ) class MockContainers: class MockContainer: def reload(self): pass def logs(self, **kwargs): return [] def __init__(self): pass def get(self, *args): pass def run(self, *args, **kwargs): assert '--native' in args[1] return MockContainers.MockContainer() class MockClient: def __init__(self, *args, **kwargs): pass def close(self): pass def version(self): return {'Version': '20.0.1'} @property def networks(self): return {'bridge': None} @property def containers(self): return MockContainers() @property def images(self): return {} monkeypatch.setattr(docker, 'from_env', MockClient) args = set_pea_parser().parse_args([ '--uses', 'docker://jinahub/pod', ]) _ = ContainerRuntime(args, ctrl_addr='', ready_event=multiprocessing.Event())
def test_idle_does_not_create_response(command, response_expected): args = set_pea_parser().parse_args([]) with Pea(args) as p: msg = ControlMessage(command, pod_name='fake_pod') with zmq.Context().socket(zmq.PAIR) as socket: socket.connect(f'tcp://localhost:{p.args.port_ctrl}') socket.send_multipart(msg.dump()) assert socket.poll(timeout=1000) == response_expected
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)
async def _create(pea: 'PeaModel'): """ .. #noqa: DAR101 .. #noqa: DAR201""" try: args = ArgNamespace.kwargs2namespace(pea.dict(), set_pea_parser()) return store.add(args) except Exception as ex: raise PartialDaemon400Exception from ex