def test_non_blocking_gateway(shards, expected_response):
class FastSlowExecutor(Executor):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
@requests(on=['/custom'])
def encode(self, docs: DocumentArray, *args, **kwargs):
assert len(docs) == 1
if docs[0].text == 'slow':
time.sleep(2)
response = []
def fill_responses(resp):
assert len(resp.data.docs) == 1
response.append(resp.data.docs[0].text)
data = DocumentArray([Document(text='slow'), Document(text='fast')])
f = Flow().add(uses=FastSlowExecutor, shards=shards, polling='ANY')
with f:
f.post(on='/custom',
inputs=data,
request_size=1,
on_done=fill_responses)
assert response == expected_response
def index(data_set, num_docs, request_size):
flow = Flow().load_config('flows/flow-index.yml')
with flow:
flow.post(on='/index',
inputs=input_index_data(num_docs, request_size, data_set),
request_size=request_size,
show_progress=True)
def test_volumes_in_flow(tmpdir, source, destination, workspace,
filewriter_exec_docker_image_built):
with mock.patch.dict(
os.environ,
{'JINA_DEFAULT_WORKSPACE_BASE': str(os.path.join(tmpdir, 'default'))},
):
if source: # test manually set volume and workspace
source = os.path.join(tmpdir, source)
volumes = [str(source) + ':' + destination]
else: # test auto volume and workspace
volumes = None
source = os.path.join(tmpdir, 'default')
f = Flow().add(uses='docker://filewriter-exec',
volumes=volumes,
workspace=workspace)
with f:
f.post(inputs=[Document()], on='/foo')
assert os.path.exists(source)
found_output_file = False # workspace has random element, so we search for it
for cur_path, dirs, files in os.walk(source):
if 'out.txt' in files:
with open(os.path.join(cur_path, 'out.txt'), 'r') as f:
if f.read() == 'Filewriter was here':
found_output_file = True
assert found_output_file
def index(workspace: Path, data_dir: Path, flow: Flow):
if workspace.exists():
shutil.rmtree(workspace)
with flow:
flow.post(
'/index', inputs=create_docs(os.path.join(data_dir, 'index', '*.mp3'))
)
def index(num_docs: int):
num_docs = min(
num_docs,
len(glob(os.path.join(os.getcwd(), IMAGE_SRC), recursive=True)))
flow = Flow(workspace="workspace")\
.add(uses={"jtype": "ImageCrafter",
"with": {"target_size": 96,
"img_mean": [0.485, 0.456, 0.406],
"img_std": [0.229, 0.224, 0.225]}}) \
.add(uses=BigTransferEncoder) \
.add(uses={"jtype": "EmbeddingIndexer",
"with": {"index_file_name": "image.json"},
"metas": {"name": "vec_idx"}},
name="vec_idx") \
.add(uses={"jtype": "KeyValueIndexer",
"metas": {"name": "kv_idx"}},
name="kv_idx",
needs="gateway") \
.add(name="join_all",
needs=["kv_idx", "vec_idx"],
read_only="true")
with flow:
document_generator = from_files(IMAGE_SRC, size=num_docs)
flow.post(on='/index',
inputs=DocumentArray(document_generator),
request_size=64,
read_mode='rb')
def test_target_peapod(mocker):
class Foo(Executor):
@requests(on='/hello')
def foo(self, **kwargs):
pass
class Bar(Executor):
@requests(on='/bye')
def bar(self, **kwargs):
pass
f = Flow().add(name='p0', uses=Foo).add(name='p1', uses=Bar)
with f:
success_mock = mocker.Mock()
fail_mock = mocker.Mock()
f.post(
'/hello',
target_peapod='p0',
inputs=Document(),
on_done=success_mock,
on_error=fail_mock,
)
success_mock.assert_called()
fail_mock.assert_not_called()
success_mock = mocker.Mock()
fail_mock = mocker.Mock()
f.post('/hello', inputs=Document(), on_done=success_mock, on_error=fail_mock)
success_mock.assert_called()
fail_mock.assert_not_called()
def test_cache_validate_remote_executor():
from .cache_validator import CacheValidator
workspace_id = random_identity()
# 1st Executor in remote workspace should download the file.
f = Flow().add(
uses=CacheValidator,
host='localhost:8000',
py_modules='cache_validator.py',
upload_files=cur_dir,
workspace_id=workspace_id,
)
with f:
response = f.post(on='/',
inputs=[Document()],
show_progress=True,
return_results=True)
assert not response[0].data.docs[0].tags['exists']
# 2nd Executor in remote workspace should be able to access the file.
f = Flow().add(
uses=CacheValidator,
host='localhost:8000',
py_modules='cache_validator.py',
upload_files=cur_dir,
workspace_id=workspace_id,
)
with f:
response = f.post(on='/',
inputs=[Document()],
show_progress=True,
return_results=True)
assert response[0].data.docs[0].tags['exists']
def test_sparse_pipeline(mocker, docs_to_index):
def validate(response):
assert len(response.docs) == 1
for doc in response.docs:
assert len(doc.matches) == TOP_K
for i, match in enumerate(doc.matches):
assert match.id == docs_to_index[i].id
assert isinstance(match.embedding, sparse.coo_matrix)
f = Flow().add(uses=DummyCSRSparseIndexEncoder)
mock = mocker.Mock()
error_mock = mocker.Mock()
with f:
f.post(
on='/index',
inputs=docs_to_index,
on_error=error_mock,
)
f.post(
on='/search',
inputs=docs_to_index[0],
parameters={'top_k': TOP_K},
on_done=mock,
on_error=error_mock,
)
mock.assert_called_once()
validate_callback(mock, validate)
error_mock.assert_not_called()
def test_func_simple_routing():
class MyExecutor(Executor):
@requests(on='/search')
def foo(self, **kwargs):
for j in ('docs', 'groundtruths', 'parameters'):
assert j in kwargs
assert len(kwargs['docs']) == 3
assert len(kwargs['groundtruths']) == 3
assert kwargs['parameters']['hello'] == 'world'
assert kwargs['parameters']['topk'] == 10
kwargs['docs'][0].tags['hello'] = 'world'
f = Flow().add(uses=MyExecutor)
with f:
results = f.post(
on='/search',
inputs=[(Document(), Document()) for _ in range(3)],
parameters={'hello': 'world', 'topk': 10},
return_results=True,
)
assert results[0].status.code == 0
assert results[0].data.docs[0].tags['hello'] == 'world'
with f:
results = f.post(
on='/random',
inputs=[Document() for _ in range(3)],
parameters={'hello': 'world', 'topk': 10},
return_results=True,
)
assert results[0].status.code == 0
def index(num_docs):
flow = Flow().load_config('flows/flow.yml')
data_path = os.path.join(os.path.dirname(__file__),
os.environ.get('JINA_DATA_FILE', None))
with flow:
flow.post(on='/index',
inputs=input_generator(num_docs, data_path),
show_progress=True)
def test_flow1(inspect, protocol, temp_folder):
f = Flow(protocol=protocol, inspect=inspect).add(
uses=DummyEvaluator1,
env={'TEST_EVAL_FLOW_TMPDIR': os.environ.get('TEST_EVAL_FLOW_TMPDIR')},
)
with f:
f.post(on='/index', inputs=docs)
def test_healthcheck_logs_websocket_with_env(capfd, health_check_env):
f = Flow(protocol='websocket', port=12345).add()
with f:
f.post('/', inputs=DocumentArray.empty())
req.get('http://localhost:12345/')
out, _ = capfd.readouterr()
assert '"GET / HTTP/1.1" 200 OK' not in out
def test_empty_post_request(mocker):
f = Flow().add(uses=MyExecutor, parallel=2, polling='ALL')
with f:
on_error_mock = mocker.Mock()
on_done_mock = mocker.Mock()
f.post('', on_error=on_error_mock, on_done=on_done_mock)
on_error_mock.assert_not_called()
on_done_mock.assert_called_once()
def test_flow_empty_data_request(mocker):
f = Flow().add(uses=MyExec)
mock = mocker.Mock()
with f:
f.post('/hello', parameters={'hello': 'world'}, on_done=mock)
mock.assert_called()
def test_flow_with_one_container_pod_shards(docker_image_built):
f = Flow().add(name='dummyEncoder1', shards=2, uses=f'docker://{img_name}')
with f:
pod = f._pod_nodes['dummyEncoder1']
assert pod.args.shards == pod.args.parallel == 2
for idx, shard in enumerate(pod.shards):
assert shard.args.pea_id == shard.args.shard_id == idx
assert shard.args.shards == shard.args.parallel == 2
f.post(on='/index', inputs=random_docs(10))
def test_load_yaml_route(mocker, req_endpoint, doc_text):
f = Flow().add(uses=y)
mock = mocker.Mock()
def validate(req):
mock()
assert req.docs[0].text == doc_text
with f:
f.post(req_endpoint, Document(), on_done=validate)
def test_flow_debug_endpoints():
f1 = Flow(protocol='http', no_debug_endpoints=True,
no_crud_endpoints=True).add(uses=MyExec)
with pytest.raises(BadClient):
with f1:
f1.post('/foo')
f2 = Flow(protocol='http', no_crud_endpoints=True).add(uses=MyExec)
with f2:
f2.post('/foo')
def test_dealer_routing(mocker):
f = Flow().add(parallel=3)
mock = mocker.Mock()
with f:
f.post(
on='/some_endpoint',
inputs=[Document() for _ in range(100)],
request_size=2,
on_done=mock,
)
mock.assert_called()
def test_flow_with_replica_container_ext_yaml(docker_image_built):
f = Flow().add(
name='dummyEncoder3',
uses=f'docker://{img_name}',
shards=3,
entrypoint='jina pea',
)
with f:
f.post(on='/index', inputs=random_docs(10))
f.post(on='/index', inputs=random_docs(10))
f.post(on='/index', inputs=random_docs(10))
def test_flow_change_parameters():
class MyExec(Executor):
@requests
def foo(self, **kwargs):
return {'a': 1}
f = Flow().add(uses=MyExec)
with f:
r = f.post('/', parameters={'a': 2}, return_results=True)
assert r[0].parameters['a'] == 1.0
r = f.post('/', parameters={}, return_results=True)
assert r[0].parameters['a'] == 1.0
def test_flow_default_polling_endpoints(polling):
f = Flow().add(uses=DynamicPollingExecutorDefaultNames,
shards=2,
polling=polling)
with f:
docs_index = f.post(on='/index', inputs=[Document(text='1')])
docs_search = f.post(on='/search', inputs=[Document(text='1')])
docs_custom = f.post(on='/custom', inputs=[Document(text='1')])
assert len(docs_index) == 2
assert len(docs_search) == 3
assert len(docs_custom) == 3 if polling == 'all' else 2
def test_change_gateway(protocol, changeto_protocol, mocker):
f = Flow(protocol=protocol).add().add().add(needs='pod1').needs_all()
with f:
mock = mocker.Mock()
f.post('/', random_docs(10), on_done=mock)
mock.assert_called()
mock = mocker.Mock()
f.protocol = changeto_protocol
f.post('/', random_docs(10), on_done=mock)
mock.assert_called()
def test_non_blocking_gateway(parallel, expected_response):
response = []
def fill_responses(resp):
assert len(resp.data.docs) == 1
response.append(resp.data.docs[0].text)
data = DocumentArray([Document(text='slow'), Document(text='fast')])
f = Flow().add(uses=FastSlowExecutor, parallel=parallel)
with f:
f.post(on='/search', inputs=data, request_size=1, on_done=fill_responses)
assert response == expected_response
def test_flow_default_custom_polling_endpoints(polling):
custom_polling_config = {'/custom': 'ALL', '/search': 'ANY', '*': polling}
f = Flow().add(
uses=DynamicPollingExecutorDefaultNames,
shards=2,
polling=custom_polling_config,
)
with f:
docs_index = f.post(on='/index', inputs=[Document(text='1')])
docs_search = f.post(on='/search', inputs=[Document(text='1')])
docs_custom = f.post(on='/custom', inputs=[Document(text='1')])
assert len(docs_index) == 2
assert len(docs_search) == 2
assert len(docs_custom) == 3
def test_func_return_():
class MyExecutor(Executor):
@requests
def foo(self, **kwargs):
return DocumentArray([Document(), Document()])
f = Flow().add(uses=MyExecutor)
with f:
f.post(
on='/some_endpoint',
inputs=[Document() for _ in range(3)],
parameters={'hello': 'world', 'topk': 10},
on_done=print,
)
def test_flow_empty_data_request(mocker):
from jina import Executor, requests
class MyExec(Executor):
@requests
def foo(self, parameters, **kwargs):
assert parameters['hello'] == 'world'
f = Flow().add(uses=MyExec)
mock = mocker.Mock()
with f:
f.post('/hello', parameters={'hello': 'world'}, on_done=mock)
mock.assert_called()
def test_func_default_routing():
class MyExecutor(Executor):
@requests
def foo(self, **kwargs):
for j in ('docs', 'groundtruths', 'parameters'):
assert j in kwargs
assert len(kwargs['docs']) == 3
f = Flow().add(uses=MyExecutor)
with f:
f.post(
on='/some_endpoint',
inputs=[Document() for _ in range(3)],
parameters={'hello': 'world', 'topk': 10},
)
def test_change_gateway(restful, changeto_gateway, mocker):
f = Flow(restful=restful).add().add().add(needs='pod1').needs_all()
with f:
mock = mocker.Mock()
f.post('', random_docs(10), on_done=mock)
mock.assert_called()
mock = mocker.Mock()
if changeto_gateway == 'RESTGateway':
f.use_rest_gateway()
if changeto_gateway == 'GRPCGateway':
f.use_grpc_gateway()
f.post('', random_docs(10), on_done=mock)
mock.assert_called()
def test_flow_load_executor_yaml_extra_search_paths():
f = Flow(extra_search_paths=[os.path.join(cur_dir, 'executor')]).add(
uses='config.yml'
)
with f:
r = f.post('/', inputs=Document(), return_results=True)
assert r[0].docs[0].text == 'done'
def run():
f = Flow().add(uses='executor1/config.yml')
# or load from Flow YAML
# f = Flow.load_config('flow.yml')
with f:
da = f.post('/', DocumentArray.empty(2))
print(da.texts)
