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)