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
