def test_uvicorn_ssl(cert_pem, key_pem, runtime_cls): args = set_gateway_parser().parse_args([ '--uvicorn-kwargs', f'ssl_certfile: {cert_pem}', f'ssl_keyfile: {key_pem}', 'ssl_keyfile_password: abcd', ]) with runtime_cls(args) as r: pass
def test_gateway_pod(runtime, protocol, runtime_cls): args = set_gateway_parser().parse_args( ['--runtime-backend', runtime, '--protocol', protocol] ) with Pod(args) as p: assert len(p.all_args) == 1 assert p.all_args[0].runtime_cls == runtime_cls Pod(args).start().close()
def _create_gateway_pod(graph_description, pod_addresses, port): return Pod(set_gateway_parser().parse_args([ '--graph-description', graph_description, '--deployments-addresses', pod_addresses, '--port', str(port), ]))
def test_custom_swagger(p): args = set_gateway_parser().parse_args(p) logger = JinaLogger('') app = get_fastapi_app(args, logger) # The TestClient is needed here as a context manager to generate the shutdown event correctly # otherwise the app can hang as it is not cleaned up correctly # see https://fastapi.tiangolo.com/advanced/testing-events/ with TestClient(app) as client: assert any('/docs' in r.path for r in app.routes) assert any('/openapi.json' in r.path for r in app.routes)
def _create_gateway_deployment(graph_description, deployments_addresses, port_expose): return Deployment(set_gateway_parser().parse_args([ '--graph-description', graph_description, '--deployments-addresses', deployments_addresses, '--port-expose', str(port_expose), ]))
def run(): args = set_gateway_parser().parse_args([ '--protocol', protocol, '--graph-description', '{}', '--deployments-addresses', '{}', ]) gateway(args)
def test_uvicorn_ssl_wrong_password(cert_pem, key_pem, runtime_cls): args = set_gateway_parser().parse_args([ '--uvicorn-kwargs', f'ssl_certfile: {cert_pem}', f'ssl_keyfile: {key_pem}', 'ssl_keyfile_password: abcde', ]) with pytest.raises(ssl.SSLError): with runtime_cls(args) as r: pass
def _create_gateway_runtime(graph_description, pod_addresses, port): with GRPCGatewayRuntime(set_gateway_parser().parse_args([ '--graph-description', graph_description, '--deployments-addresses', pod_addresses, '--port', str(port), ])) as runtime: runtime.run_forever()
def _create_gateway_pod(l_or_r, graph_description, deployments_addresses, port_expose): args = set_gateway_parser().parse_args([]) if l_or_r == 'remote': args.host = HOST args.port_jinad = PORT args.graph_description = graph_description args.deployments_addresses = deployments_addresses args.port_expose = port_expose args.runtime_cls = 'GRPCGatewayRuntime' return PodFactory.build_pod(args)
def test_gateway_pod(runtime, restful, runtime_cls): args = set_gateway_parser().parse_args( ['--runtime-backend', runtime, '--runtime-cls', runtime_cls] + (['--restful'] if restful else [])) with Pod(args) as p: assert len(p.all_args) == 1 if restful: assert p.all_args[0].runtime_cls == 'RESTRuntime' else: assert p.all_args[0].runtime_cls == 'GRPCRuntime' Pod(args).start().close()
def _create_runtime(): with GRPCGatewayRuntime(set_gateway_parser().parse_args([ '--port-expose', str(port_expose), '--graph-description', '{"start-gateway": ["deployment0"], "deployment0": ["end-gateway"]}', '--deployments-addresses', '{"deployment0": ["0.0.0.0:' + f'{deployment0_port}' + '", "0.0.0.0:' + f'{deployment1_port}' + '"]}', ])) as runtime: runtime.run_forever()
def test_gateway_args(protocol, expected): args = set_gateway_parser().parse_args([ '--host', 'jina-custom-gateway', '--port-expose', '23456', '--protocol', protocol, ]) p = Pea(args) assert p.runtime_cls.__name__ == expected
async def test_concurrent_requests(): args = ArgNamespace.kwargs2namespace({}, set_gateway_parser()) mock_zmqlet = ZmqletMock() streamer = ZmqGatewayStreamer(args, mock_zmqlet) request = _generate_request() response = streamer.stream(iter([request])) async for r in response: assert r.proto == request
def test_gateway_runtimes(protocol, expected): args = set_gateway_parser().parse_args([ '--graph-description', '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}', '--deployments-addresses', '{"pod0": ["0.0.0.0:1234"]}', '--protocol', protocol, ]) with Pod(args) as p: assert p.runtime_cls.__name__ == expected
async def test_concurrent_requests(): args = ArgNamespace.kwargs2namespace({}, set_gateway_parser()) mock_zmqlet = ZmqletMock() servicer = PrefetchCaller(args, mock_zmqlet) request = _generate_request() response = servicer.send(iter([request])) async for r in response: assert r.proto == request await servicer.close()
def test_failing_gateway_runtimes(protocol, expected): args = set_gateway_parser().parse_args([ '--graph-description', '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}', '--deployments-addresses', '{_INVALIDJSONINTENTIONALLY_pod0": ["0.0.0.0:1234"]}', '--protocol', protocol, ]) with pytest.raises(RuntimeFailToStart): with Pod(args): pass
def _create_gateway_pod(graph_description, pod_addresses, port_expose): return Pod( set_gateway_parser().parse_args( [ '--graph-description', graph_description, '--deployments-addresses', pod_addresses, '--port-expose', str(port_expose), '--noblock-on-start', ] ) )
def test_gateway_pod(protocol, runtime_cls, graph_description): args = set_gateway_parser().parse_args([ '--graph-description', graph_description, '--deployments-addresses', '{"pod0": ["0.0.0.0:1234"]}', '--protocol', protocol, ]) with Deployment(args) as p: assert len(p.all_args) == 1 assert p.all_args[0].runtime_cls == runtime_cls Deployment(args).start().close()
def dump(self, data: 'Flow') -> Dict: """ :param data: versioned flow object :return: the dictionary given a versioned flow object """ r = {} if data._version: r['version'] = data._version # to maintain order - version -> with -> executors r['with'] = {} if data._kwargs: r['with'].update(data._kwargs) if data._common_kwargs: r['with'].update(data._common_kwargs) if data._deployment_nodes: r['executors'] = [] last_name = 'gateway' for k, v in data._deployment_nodes.items(): kwargs = {} # only add "needs" when the value is not the last deployment name if list(v.needs) != [last_name]: kwargs = {'needs': list(v.needs)} # get nondefault kwargs parser = set_deployment_parser() if v.role == DeploymentRoleType.GATEWAY: parser = set_gateway_parser() non_default_kw = ArgNamespace.get_non_defaults_args(v.args, parser) kwargs.update(non_default_kw) for t in _get_taboo(parser): if t in kwargs: kwargs.pop(t) if k == 'gateway': if 'JINA_FULL_CLI' in os.environ: r['with'].update(kwargs) else: continue else: last_name = kwargs['name'] r['executors'].append(kwargs) return r
def _create_gateway_pod(graph_description, pod_addresses, port, protocol='grpc'): return Pod(set_gateway_parser().parse_args([ '--graph-description', graph_description, '--deployments-addresses', pod_addresses, '--port', str(port), '--noblock-on-start', '--protocol', protocol, ]))
def test_gateway_ready(port_expose, route, status_code): p = set_gateway_parser().parse_args([ '--port-expose', str(port_expose), '--protocol', 'http', '--graph-description', '{}', '--deployments-addresses', '{}', ]) with PodFactory.build_pod(p): time.sleep(0.5) a = requests.get(f'http://localhost:{p.port_expose}{route}') assert a.status_code == status_code
def _create_gqlgateway_runtime(graph_description, pod_addresses, port): with HTTPGatewayRuntime( set_gateway_parser().parse_args( [ '--graph-description', graph_description, '--deployments-addresses', pod_addresses, '--port', str(port), '--expose-graphql-endpoint', ] ) ) as runtime: runtime.run_forever()
def run(*args, **kwargs): runtime_cls = None print(f' args {args}') runtime_args = set_gateway_parser().parse_args(args) print(f' protocol {runtime_args.protocol}') if runtime_args.protocol == GatewayProtocolType.GRPC: runtime_cls = GRPCGatewayRuntime elif runtime_args.protocol == GatewayProtocolType.HTTP: runtime_cls = HTTPGatewayRuntime elif runtime_args.protocol == GatewayProtocolType.WEBSOCKET: runtime_cls = WebSocketGatewayRuntime print(f' runtime_cls {runtime_cls}') with runtime_cls(runtime_args) as runtime: print(f' Lets run forever') runtime.run_forever()
def get_k8s_pod( pod_name: str, namespace: str, shards: str = None, replicas: str = None, needs: Optional[Set[str]] = None, uses_before=None, uses_after=None, port_expose=None, ): parameter_list = ['--name', pod_name, '--k8s-namespace', namespace] if shards: parameter_list.extend( [ '--shards', str(shards), ] ) if replicas: parameter_list.extend( [ '--replicas', str(replicas), ] ) if port_expose: parameter_list.extend( [ '--port-expose', str(port_expose), ] ) if uses_before: parameter_list.extend( [ '--uses-before', uses_before, ] ) if uses_after: parameter_list.extend(['--uses-after', uses_after]) parameter_list.append('--noblock-on-start') parser = set_gateway_parser() if pod_name == 'gateway' else set_pod_parser() args = parser.parse_args(parameter_list) pod = K8sPod(args, needs) return pod
def get_gateway_yamls( self, ) -> List[Dict]: import os image_name = os.getenv( 'JINA_GATEWAY_IMAGE', f'jinaai/jina:{self.version}-py38-standard' ) cargs = copy.copy(self.deployment_args) cargs.deployments_addresses = self.k8s_deployments_addresses from jina.helper import ArgNamespace from jina.parsers import set_gateway_parser taboo = { 'uses_with', 'uses_metas', 'volumes', 'uses_before', 'uses_after', 'workspace', 'workspace_id', 'upload_files', 'noblock_on_start', 'env', } non_defaults = ArgNamespace.get_non_defaults_args( cargs, set_gateway_parser(), taboo=taboo ) _args = ArgNamespace.kwargs2list(non_defaults) container_args = ['gateway'] + _args return kubernetes_deployment.get_deployment_yamls( self.dns_name, namespace=self.k8s_namespace, image_name=image_name, container_cmd='["jina"]', container_args=f'{container_args}', replicas=1, pull_policy='IfNotPresent', jina_deployment_name='gateway', pod_type=self.pod_type, port=self.common_args.port, env=cargs.env, monitoring=self.common_args.monitoring, port_monitoring=self.common_args.port_monitoring, )
def get_gateway_yamls(self, ) -> List[Dict]: import os test_pip = os.getenv('JINA_K8S_USE_TEST_PIP') is not None image_name = ('jinaai/jina:test-pip' if test_pip else f'jinaai/jina:{self.version}-py38-standard') cargs = copy.copy(self.deployment_args) cargs.env = None cargs.deployments_addresses = self.k8s_deployments_addresses from jina.helper import ArgNamespace from jina.parsers import set_gateway_parser taboo = { 'uses_with', 'uses_metas', 'volumes', 'uses_before', 'uses_after', 'workspace', 'workspace_id', 'upload_files', 'noblock_on_start', } non_defaults = ArgNamespace.get_non_defaults_args( cargs, set_gateway_parser(), taboo=taboo) _args = ArgNamespace.kwargs2list(non_defaults) container_args = ['gateway'] + _args if not cargs.k8s_connection_pool: container_args.append('--k8s-disable-connection-pool') return kubernetes_deployment.get_deployment_yamls( self.dns_name, namespace=self.k8s_namespace, image_name=image_name, container_cmd='["jina"]', container_args=f'{container_args}', replicas=1, pull_policy='IfNotPresent', jina_deployment_name='gateway', pod_type=self.pod_type, port_expose=self.common_args.port_expose, env=cargs.env, )
def _create_gateway_runtime(graph_description, pod_addresses, port, protocol='grpc'): if protocol == 'http': gateway_runtime = HTTPGatewayRuntime elif protocol == 'websocket': gateway_runtime = WebSocketGatewayRuntime else: gateway_runtime = GRPCGatewayRuntime with gateway_runtime(set_gateway_parser().parse_args([ '--graph-description', graph_description, '--deployments-addresses', pod_addresses, '--port', str(port), ])) as runtime: runtime.run_forever()
def test_index_remote_rpi(test_workspace): f_args = set_gateway_parser().parse_args(['--host', '0.0.0.0']) def start_gateway(): with Pod(f_args): time.sleep(3) t = mp.Process(target=start_gateway) t.daemon = True t.start() f = Flow(optimize_level=FlowOptimizeLevel.IGNORE_GATEWAY).add( uses=os.path.join(cur_dir, 'yaml/test-index-remote.yml'), parallel=3, host='0.0.0.0', port_expose=random_port()) with f: f.index(input_fn=random_docs(1000))
def create_runtime( graph_dict: Dict, protocol: str, port: int, call_counts=None, monkeypatch=None ): import json graph_description = json.dumps(graph_dict) runtime_cls = None if call_counts: def decompress(self): call_counts.put_nowait('called') from jina.proto import jina_pb2 self._pb_body = jina_pb2.DataRequestProto() self._pb_body.ParseFromString(self.buffer) self.buffer = None monkeypatch.setattr( DataRequest, '_decompress', decompress, ) if protocol == 'grpc': runtime_cls = GRPCGatewayRuntime elif protocol == 'http': runtime_cls = HTTPGatewayRuntime elif protocol == 'websocket': runtime_cls = WebSocketGatewayRuntime with runtime_cls( set_gateway_parser().parse_args( [ '--port', f'{port}', '--graph-description', f'{graph_description}', '--deployments-addresses', '{}', ] ) ) as runtime: runtime.run_forever()
def process_wrapper(): monkeypatch.setattr( networking.GrpcConnectionPool, 'send_requests_once', DummyMockConnectionPool.send_requests_once, ) port_in = random_port() with GRPCGatewayRuntime(set_gateway_parser().parse_args([ '--port-expose', f'{port_in}', '--graph-description', f'{json.dumps(complete_graph_dict)}', '--deployments-addresses', '{}', ])) as runtime: async def _test(): responses = [] req = request_generator( '/', DocumentArray([Document(text='client0-Request')])) async for resp in runtime.streamer.Call(request_iterator=req): responses.append(resp) return responses responses = asyncio.run(_test()) assert len(responses) > 0 assert len(responses[0].docs) == 1 deployment2_path = ( f'client0-Request-client0-deployment0-client0-deployment2-client0-deployment3-client0-merger-client0-deployment_last' == responses[0].docs[0].text) deployment4_path = ( f'client0-Request-client0-deployment4-client0-deployment5-client0-merger-client0-deployment_last' == responses[0].docs[0].text) assert ( f'client0-Request-client0-deployment0-client0-deployment1-client0-merger-client0-deployment_last' == responses[0].docs[0].text or deployment2_path or deployment4_path)