Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
feat(gateway): https support to uvicorn based servers (#3444)
* feat(gateway): https support to uvicorn based servers

* style: fix overload and cli autocomplete

Co-authored-by: Jina Dev Bot <dev-bot@jina.ai>
  • Loading branch information
deepankarm and jina-bot committed Sep 20, 2021
1 parent 4a5eee9 commit cb46192
Show file tree
Hide file tree
Showing 10 changed files with 244 additions and 7 deletions.
11 changes: 10 additions & 1 deletion cli/autocomplete.py
Expand Up @@ -172,6 +172,7 @@
'--no-debug-endpoints',
'--no-crud-endpoints',
'--expose-endpoints',
'--uvicorn-kwargs',
'--compress',
'--compress-min-bytes',
'--compress-min-ratio',
Expand Down Expand Up @@ -347,7 +348,15 @@
'--k8s-namespace',
'--k8s-custom-resource-dir',
],
'client': ['--help', '--host', '--proxy', '--port', '--asyncio', '--protocol'],
'client': [
'--help',
'--host',
'--proxy',
'--port',
'--https',
'--asyncio',
'--protocol',
],
'export-api': ['--help', '--yaml-path', '--json-path', '--schema-path'],
},
}
6 changes: 4 additions & 2 deletions jina/clients/__init__.py
@@ -1,12 +1,12 @@
"""Module wrapping the Client of Jina."""
import argparse
from typing import overload, Optional, Union
from typing import overload, Optional, Union, TYPE_CHECKING

__all__ = ['Client']

from ..enums import GatewayProtocolType

if False:
if TYPE_CHECKING:
from .grpc import GRPCClient, AsyncGRPCClient
from .websocket import WebSocketClient, AsyncWebSocketClient
from .http import HTTPClient, AsyncHTTPClient
Expand All @@ -18,6 +18,7 @@ def Client(
*,
asyncio: Optional[bool] = False,
host: Optional[str] = '0.0.0.0',
https: Optional[bool] = False,
port: Optional[int] = None,
protocol: Optional[str] = 'GRPC',
proxy: Optional[bool] = False,
Expand All @@ -34,6 +35,7 @@ def Client(
:param asyncio: If set, then the input and output of this Client work in an asynchronous manner.
:param host: The host address of the runtime, by default it is 0.0.0.0.
:param https: If set, connect to gateway using https
:param port: The port of the Gateway, which the client should connect to.
:param protocol: Communication protocol between server and client.
:param proxy: If set, respect the http_proxy and https_proxy environment variables. otherwise, it will unset these proxy variables before start. gRPC seems to prefer no proxy
Expand Down
3 changes: 2 additions & 1 deletion jina/clients/base/http.py
Expand Up @@ -60,7 +60,8 @@ async def _get_results(

try:
cm1 = ProgressBar() if self.show_progress else nullcontext()
url = f'http://{self.args.host}:{self.args.port}/post'
proto = 'https' if self.args.https else 'http'
url = f'{proto}://{self.args.host}:{self.args.port}/post'

with cm1 as p_bar:
all_responses = []
Expand Down
4 changes: 3 additions & 1 deletion jina/clients/base/websocket.py
Expand Up @@ -39,10 +39,12 @@ async def _get_results(
# setting `max_size` as None to avoid connection closure due to size of message
# https://websockets.readthedocs.io/en/stable/api.html?highlight=1009#module-websockets.protocol

proto = 'wss' if self.args.https else 'ws'
async with websockets.connect(
f'ws://{self.args.host}:{self.args.port}/',
f'{proto}://{self.args.host}:{self.args.port}/',
max_size=None,
ping_interval=None,
ssl=self.args.https or None,
) as websocket:
# To enable websockets debug logs
# https://websockets.readthedocs.io/en/stable/cheatsheet.html#debugging
Expand Down
6 changes: 6 additions & 0 deletions jina/flow/base.py
Expand Up @@ -85,6 +85,7 @@ def __init__(
*,
asyncio: Optional[bool] = False,
host: Optional[str] = '0.0.0.0',
https: Optional[bool] = False,
port: Optional[int] = None,
protocol: Optional[str] = 'GRPC',
proxy: Optional[bool] = False,
Expand All @@ -94,6 +95,7 @@ def __init__(
:param asyncio: If set, then the input and output of this Client work in an asynchronous manner.
:param host: The host address of the runtime, by default it is 0.0.0.0.
:param https: If set, connect to gateway using https
:param port: The port of the Gateway, which the client should connect to.
:param protocol: Communication protocol between server and client.
:param proxy: If set, respect the http_proxy and https_proxy environment variables. otherwise, it will unset these proxy variables before start. gRPC seems to prefer no proxy
Expand Down Expand Up @@ -160,6 +162,7 @@ def __init__(
uses_metas: Optional[dict] = None,
uses_requests: Optional[dict] = None,
uses_with: Optional[dict] = None,
uvicorn_kwargs: Optional[dict] = None,
workspace: Optional[str] = None,
zmq_identity: Optional[str] = None,
**kwargs,
Expand Down Expand Up @@ -250,6 +253,9 @@ def __init__(
:param uses_metas: Dictionary of keyword arguments that will override the `metas` configuration in `uses`
:param uses_requests: Dictionary of keyword arguments that will override the `requests` configuration in `uses`
:param uses_with: Dictionary of keyword arguments that will override the `with` configuration in `uses`
:param uvicorn_kwargs: Dictionary of kwargs arguments that will be passed to Uvicorn server when starting the server
More details can be found in Uvicorn docs: https://www.uvicorn.org/settings/
:param workspace: The working directory for any IO operations in this object. If not set, then derive from its parent `workspace`.
:param zmq_identity: The identity of a ZMQRuntime. It is used for unique socket identification towards other ZMQRuntimes.
Expand Down
22 changes: 21 additions & 1 deletion jina/parsers/peapods/runtimes/remote.py
@@ -1,5 +1,5 @@
"""Argparser module for remote runtime"""
from ...helper import add_arg_group
from ...helper import KVAppendAction, add_arg_group
from .... import __default_host__
from .... import helper
from ....enums import CompressAlgo
Expand Down Expand Up @@ -35,6 +35,13 @@ def mixin_client_gateway_parser(parser):
help='The port of the Gateway, which the client should connect to.',
)

gp.add_argument(
'--https',
action='store_true',
default=False,
help='If set, connect to gateway using https',
)


def mixin_gateway_parser(parser):
"""Add the options for remote expose at the Gateway
Expand Down Expand Up @@ -134,6 +141,19 @@ def mixin_http_gateway_parser(parser=None):
''',
)

gp.add_argument(
'--uvicorn-kwargs',
action=KVAppendAction,
metavar='KEY: VALUE',
nargs='*',
help='''
Dictionary of kwargs arguments that will be passed to Uvicorn server when starting the server
More details can be found in Uvicorn docs: https://www.uvicorn.org/settings/
''',
)


def mixin_prefetch_parser(parser=None):
"""Add the options for prefetching
Expand Down
2 changes: 2 additions & 0 deletions jina/peapods/runtimes/gateway/http/__init__.py
Expand Up @@ -50,12 +50,14 @@ async def serve(self, **kwargs):

from .....helper import extend_rest_interface

uvicorn_kwargs = self.args.uvicorn_kwargs or {}
self._server = UviServer(
config=Config(
app=extend_rest_interface(get_fastapi_app(self.args, self.logger)),
host=__default_host__,
port=self.args.port_expose,
log_level=os.getenv('JINA_LOG_LEVEL', 'error').lower(),
**uvicorn_kwargs
)
)
await self._server.setup()
Expand Down
2 changes: 2 additions & 0 deletions jina/peapods/runtimes/gateway/websocket/__init__.py
Expand Up @@ -50,13 +50,15 @@ async def serve(self, **kwargs):

from .....helper import extend_rest_interface

uvicorn_kwargs = self.args.uvicorn_kwargs or {}
self._server = UviServer(
config=Config(
app=extend_rest_interface(get_fastapi_app(self.args, self.logger)),
host=__default_host__,
port=self.args.port_expose,
ws_max_size=1024 * 1024 * 1024,
log_level=os.getenv('JINA_LOG_LEVEL', 'error').lower(),
**uvicorn_kwargs
)
)
await self._server.setup()
Expand Down
28 changes: 28 additions & 0 deletions tests/unit/parsers/peapods/runtimes/test_remote_parser.py
@@ -0,0 +1,28 @@
import argparse
from jina.parsers.peapods.runtimes.remote import mixin_http_gateway_parser


def test_runtime_parser():
parser = argparse.ArgumentParser(
epilog=f'Test', description='Test Command Line Interface'
)

mixin_http_gateway_parser(parser)

args = parser.parse_args([])
assert args.uvicorn_kwargs is None

args = parser.parse_args(['--uvicorn-kwargs', 'ssl_keyfile: /tmp/cert.pem'])
assert args.uvicorn_kwargs == {'ssl_keyfile': '/tmp/cert.pem'}

args = parser.parse_args(
[
'--uvicorn-kwargs',
'ssl_keyfile: /tmp/cert.pem',
'ssl_keyfile_password: 1234e',
]
)
assert args.uvicorn_kwargs == {
'ssl_keyfile': '/tmp/cert.pem',
'ssl_keyfile_password': '1234e',
}
167 changes: 166 additions & 1 deletion tests/unit/peapods/runtimes/gateway/http/test_app.py
@@ -1,12 +1,17 @@
import ssl
from tempfile import NamedTemporaryFile

import pytest
import requests as req
from fastapi.testclient import TestClient

from jina import Document, Client
from jina.helper import random_port
from jina import Executor, requests, Flow, DocumentArray
from jina.logging.logger import JinaLogger
from jina.parsers import set_gateway_parser
from jina.peapods.runtimes.gateway.http import get_fastapi_app
from jina.peapods.runtimes.gateway.websocket import WebSocketRuntime
from jina.peapods.runtimes.gateway.http import HTTPRuntime, get_fastapi_app


@pytest.mark.parametrize('p', [['--default-swagger-ui'], []])
Expand Down Expand Up @@ -48,3 +53,163 @@ def test_tag_update(grpc_data_requests):
assert r1.json()['data']['docs'][0]['tags'] == {'prop1': 'val'}
r2 = req.post(f'http://localhost:{PORT_EXPOSE}/index', json=d2)
assert r2.json()['data']['docs'][0]['tags'] == {'prop2': 'val'}


@pytest.fixture
def cert_pem():
"""This is the cert entry of a self-signed local cert"""
tmp = NamedTemporaryFile('w')
tmp.write(
"""-----BEGIN CERTIFICATE-----
MIIFazCCA1OgAwIBAgIUE663J9NKJE5sTDXei0ScmKE1TskwDQYJKoZIhvcNAQEL
BQAwRTELMAkGA1UEBhMCQVUxEzARBgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoM
GEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDAeFw0yMTA5MjAwOTQ4NThaFw0yMjA5
MjAwOTQ4NThaMEUxCzAJBgNVBAYTAkFVMRMwEQYDVQQIDApTb21lLVN0YXRlMSEw
HwYDVQQKDBhJbnRlcm5ldCBXaWRnaXRzIFB0eSBMdGQwggIiMA0GCSqGSIb3DQEB
AQUAA4ICDwAwggIKAoICAQC243Axri0Aafq5VsS5+w1QgSIYjhWWCi0frm/w95+O
SleiyQ2nR6Cas2YHViLPo4casch+M5d7fzxzSyezKLoM9FJ7p9rHAc08sjuIkqMt
kDApgfl4Rtco/KqgEr0HELpo6rWG8tby0Wbl82eSm93GUAyZwyuZMdr3Ag6v8ppn
JaUit1oWWs8XZdvEIoRxXQu+APNiKaWWrFjbSXay/ZxbsDrdk7Q+bHLiwYxhx3Bj
SZX9xWPjchFv+fD1pBOyq/P76VGr6B938vEj+EorqUwdiIeW3vgw2FODLg5bXMSo
YR6uZ1V2W8xGwWHpj0s1UChbaOY9thRxvtOrKeW9F4xoFoBrr6ZjkcqD/5mARJz+
Uwee/XhLE7Z5L+eyzLXcXLR2lOs8AXgCmUgAgk0NJi8IPQGZFEBuWVJ7DBO87G7p
DbKMkQ4QGB4dj7lJdHUr6v07Z+Etus7Z+cwjQWe2wdQgDV05E/zCSwWIv4AYbGXs
s1P4XXMeYxxK/74vh7k15TmIiq77A96FaxStK2PZXJjI1dB5DhoC93qCZogq4vup
r6Yk6B29whOlHsBWVL4nW6SYxEDNKyWYRRekiJlcxlw+NpZxBUdC5PwOkh4AZmnW
PpBZv/rCXC7Ow0DS9F9CbfzVynihUHLlZk2SvH8Dc4htum+guiwBMyRtNaSdD8l2
OwIDAQABo1MwUTAdBgNVHQ4EFgQUvTljFuE/DJlq0s8U3wdteIHmQbwwHwYDVR0j
BBgwFoAUvTljFuE/DJlq0s8U3wdteIHmQbwwDwYDVR0TAQH/BAUwAwEB/zANBgkq
hkiG9w0BAQsFAAOCAgEAh7yvPSX3qzWtczJNO4HGzo0xJC9vnEy01xKLOPZiGSFE
CFA15mtCnX1hYglkK8Bp7UnKMo9MwQFC4WLjBj7Ts6NamR0KpMf5M75WrCCA55UT
aWbwqFmHH47j7DYl/j1WK/nLocAzW3zyjDQaBArAls9Iw5nkeonyl+yaXh/7J45G
tNRrMyyxk1hl2C4NA3grQHycZiviNr6OQmgZ3sPexaPfAva3Zuwwv24UeliB0Lpb
opxcZJ9ojqSAok/eJCKSpywuVkxy61Iz2MKIpLA+WoRFjVGuvM5rZPSEQweWlnJT
f4GVKdfGQW9bzM27fMse/sg03z6odTn0rkxUM8TWsZR3Jg9YKbP2zgo4snU9FUMZ
RQA1A83U1T10yaeaCLBjN2psATQr7blYZhNUwYVr41C4K9+g3ghK3zhrKeaokBKQ
xo1aQZQNMyxrpe6NU+Iu9Esz4LRKaf8B4Q5vXJhf2YPqaz3FSHOFHNTiILvIEnuD
DFRwYLPkWlFLr5MYyjo8IlL/lcAjv3F3+Nx7qfvtIoLLxVON4hacYpG2uyyDGqg0
TiIvOLZ67W63nUk6h7+Pwm/8EhxTxFjguSOh0fu7GXtF75kDueBLoERr6DgcBTTg
adVnffnjz+hTFEjwXL48iGRPM142AGNOfXNp8tvPZOYjkc2prtIhGlvOu+De8tg=
-----END CERTIFICATE-----"""
)
tmp.flush()
yield tmp.name
tmp.close()


@pytest.fixture
def key_pem():
"""This is the key entry of a self-signed local cert"""
tmp = NamedTemporaryFile('w')
tmp.write(
"""-----BEGIN ENCRYPTED PRIVATE KEY-----
MIIJnDBOBgkqhkiG9w0BBQ0wQTApBgkqhkiG9w0BBQwwHAQIQZi3yv841tUCAggA
MAwGCCqGSIb3DQIJBQAwFAYIKoZIhvcNAwcECJlWkgQTVxuQBIIJSMEqOJZ6gnFe
NPER3S6xBLwRToNvspoj9gl8befzOEgLEHv5TFQ2WCznQ5+HkrufnSVwu/R8XR68
CvrIhjYqmO3TAI0t8PWeKOA2Zw3AI5Qlry3L3HouPXoj7OMBs3A99bjgdjIV7C1b
EaiMA9RsBe3jZfaqHnBX19n6pymej9JYJKmqaY2b5+jRh7bu36+0J4/TbLYH3AKX
U7nKG+cKmbYdrAx3ftHuSwTO8KgXy2fDxjqicGJ1D80RT0Pk+jUAYdZg6OWwGwiP
6qeObJikseR/FuSqYhuH5vf9ialLBVR2jmWZw7xJh+Kv4T7GZs0V/hru31cYQU+b
2IcpfBJsR9mKpilD/A62obCk0TSJDY8gehmtvhyZpGnFyh2eOESMuIPpvvsR2kdS
9j7E2luksuMATcdJMElFgM2xS0il0H1amsCtQWgK9BFDlzmOOxr6K6Hm6MS+dq20
1nvznQFt1Qd5e9hyuPV8qd/uuqA+BlnwAJds++fR8jB55ZNOfWzD5sQky+wGkJ80
CwzOQLgqKQqUyYR/+SD8dfTyOPNyBu5f1RBkb2gTOTRQgwDQUOOdwfNmb9brHj0Y
/c6zQ6UCkXYKXxjBr/O7S2yTwsC/gnN3PUEQBWCbYlIe/A0EorouOKMFwCLj1r/U
fn3H0XvuY0ahqNZBxVnuWcUsSGzMiaePIsJsWXmz0A4ufS+CqSIwGx28A/cYhC/2
yZzss3yACFCHIJeoPpOsKPcyhw8K48YyofW5L8vKI5eadhbjDyRH38Z2zSYGLw/1
YRwxsAVRimnLUiz2GND2XXBLEJFLthd8BPnQM1+CEun54UeQaVRJ8p1PyPGIWvWS
Nb4o8jBIjMavwjxjCF7WCdO7V0iThcPwCLme9AN2+MaqWC1HBVZ1QRAVAoVsEFXz
4TPjJljn3RYf8anGfRxsNX74QqatCL2+Oy1M48vqAUICaTPh6z88hpBCYAyvRiXz
S9CEs9o/TcBtcemF1AYB7tsmkYkaJJaOd3M+t/0VFuXYv4OdqiY5PtsdSFh+j1KJ
u1+Pav+jqGrJ/zJuLOC6dx2yaoT1TWRcMcUT3r7j7+yAL97D8azo7PW2F2a9zHl5
2HgS4byzHSRPsGSyKOtGQT2hz12TzgsJ0YPvUz9z4Hi62PgPfNWyukZovWEhxu0X
F1ll2xd/g4QYCqs1dsCU8IQB6xBUbLJ5noQNGN1JvAqKTDBENr28cD2r+QcZPuE3
84LbQJLWCJfwYHJ5GyWFNWyb4DjfZ3HdVHSvRVdsjYLhJDEWKNyW9EJi9hW3PujI
CEZgW3JfWkVmUj64DbmkervtpUM//J/KOdCHIaxKcjSbJziiQfI0q+hR5VS9ceaE
9AYcviAQ3few5MNt869HeHfxGfuG1FDKBSmtf5bLbtlx+RIq79bkFl/A3esD0q0t
2UccHTorNBgVDKkBLETuyCyugiOe6XEpVD+gooW39C+fWe7dxeN7uWYB+IVTfaDB
qKMrgiuUkZZmO2B0YLoDCsvnVlOlH2tvqm6DSAn8BsKU5LzIGDv4g04CCMDkHt0I
8DoUFFhjPHOwGK40gtsekFLz3DlU5c43AVcW9V8pV8A4m5+ZXWI0re6M38QzEaIM
Mtso1Mq/y8gyE+iB0Y1Tx3OY3l0FDmyAwAzCMbkhcI1OcUW37/43wi4MGk9NPaBH
t7XaiLo78jpH9Y1jC1zhgIrcllNWBzlm1Nh94ZrcDk2YZt4c49Lg0+ghO1CW1IH5
ulGjRn3z/sGrjHfGF4GNICbxODrWidXC58/dRh515BK848sFnQCCTVYR6dARhTQQ
13zEvzXX2UJHDpbE0ut1Z0A4IVfvG0ZUoZGGTx+TZFKalKyFJh7/be19gg7K+1z8
BswuwkIvRbsQaxq9BlzS11clOLPr5gu9DOAICJb8tscPa+B0PC7TgZ5JpB3Gv/GS
zdslokIN+gEGUINZFVTLOJVvactkFNO/bCM/TSdn/5LmSJ9MbkYYhpIgPs3Oz1ia
E6Xq9tacvyeOWp9rbz2LG03iMQd45slsPoGyQPOsvmZ48SipuefkWmMLA3LuB71/
IOeO/MIQ29qunr8edkEm2uV9GS+09JUVOr4N/Ir9OGmr1UPkFenEnbtSiYzSQDov
FLIMj4p1KoPcQDwHPsqj9hF47rgArJ6RWZlMo2vDA4bcTTxKugHPaitedJ2d+WJj
fs+5C6D8E8lXpb3oh3ncsFQt7LGJWBOQYaxhPzfAdX5/s9CIqHyIStEY38/Izs+F
sgC6YUOk+5j5IIik6C65YG9mcQwHvCYWynch4PSpa87qjDkP/3BQWNb5OCcwsZ24
lap/PkIXxMKHsoh3i4moQDcaKUEPF4cgzOj/+IipMu/MCizNAm8bhaS2JRKOXGIN
eU9bsw+ADHMrtiLHiEH98ifabCGadvp8B8ZkpYpcT/LtwkHjJ4x7AMFEKK1Cj92r
eaiYszKVYwuTZObGNkWta6AiIsoqU84/NFUpaGn2Qdr4FK6+YBddhlUPs+amrOZF
hy8I5qP6WqNtKmVyPHWY96OhR9JmYxlpVWYr5UzhJ+JClTnVqy++K+j91JahyCBa
1d8ey15Ibhlu3nQ5stxpnmzA/NpBCLhTFUuri/T3C91hHKJicjeZFYpumjHOPZXA
+JwUGwsIkO7n9KiA6F7IOJgJIMHE3VeO6QLdiA3QJbj47o7vwQLnMwOByKrEGIQP
yKERA6oZft29EqqNBAxgp3hDXQI7SIjHVtq1kuTmwu8o7Y5vFxG/r1D+W2V3ZAsr
atXA7La/FbQwfFvCaWPtCy+vehiKjdr2Z+4/BrkTPtkaMe+1qMc21K2rYZaIw+mh
p++zQ0j+16Y5putGcPPf8i1vQ0eMd2OljXo0bqLn5n3b0Q3arRnPpXumgXZES90I
wJCkQIiAy+AYoLROVVrefmQ4/XlWA5iizqkTDU1NThxSQjVan14O356G0HmxNsi9
RB2a0AmwuGhuYPYjUI8iKU12RMp4/rRb28xbAwSh24qQeY2a/IY4u6bGpOWdTudg
Xb3L8FmNUZVtO0QvLKa6YHUW0BTgUy4EzA9nDKDRMYIrRh3BMTr2YZ4rA5ReY1+T
lFkijOU5iJjWLTYGcCyBHQup/VrqmgxchRbbKFO5+qpDHE0e3oLbPLQ0Rw425SvN
xZ36Vrgc4hfaUiifsIiDwA==
-----END ENCRYPTED PRIVATE KEY-----"""
)
tmp.flush()
yield tmp.name
tmp.close()


@pytest.mark.parametrize('runtime_cls', [HTTPRuntime, WebSocketRuntime])
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


@pytest.mark.parametrize('runtime_cls', [HTTPRuntime, WebSocketRuntime])
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


@pytest.mark.parametrize('protocol', ['http', 'websocket'])
def test_uvicorn_ssl_with_flow(cert_pem, key_pem, protocol, monkeypatch):
with Flow(
protocol=protocol,
uvicorn_kwargs=[
f'ssl_certfile: {cert_pem}',
f'ssl_keyfile: {key_pem}',
'ssl_keyfile_password: abcd',
],
) as f:
if protocol == 'http':
Client(protocol=protocol, port=f.port_expose, https=True).index(
[Document()]
)
else:
with pytest.raises(ssl.SSLCertVerificationError) as r:
Client(protocol=protocol, port=f.port_expose, https=True).index(
[Document()]
)
assert (
'certificate verify failed: self signed certificate' in r.value.args[1]
)

0 comments on commit cb46192

Please sign in to comment.