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)
