def get_fastapi_app(
args: 'argparse.Namespace',
topology_graph: 'TopologyGraph',
connection_pool: 'GrpcConnectionPool',
logger: 'JinaLogger',
metrics_registry: Optional['CollectorRegistry'] = None,
):
"""
Get the app from FastAPI as the REST interface.
:param args: passed arguments.
:param topology_graph: topology graph that manages the logic of sending to the proper executors.
:param connection_pool: Connection Pool to handle multiple replicas and sending to different of them
:param logger: Jina logger.
:param metrics_registry: optional metrics registry for prometheus used if we need to expose metrics from the executor or from the data request handler
:return: fastapi app
"""
with ImportExtensions(required=True):
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from starlette.requests import Request
from jina.serve.runtimes.gateway.http.models import (
JinaEndpointRequestModel,
JinaRequestModel,
JinaResponseModel,
JinaStatusModel,
)
docs_url = '/docs'
app = FastAPI(
title=args.title or 'My Jina Service',
description=args.description or
'This is my awesome service. You can set `title` and `description` in your `Flow` or `Gateway` '
'to customize this text.',
version=__version__,
docs_url=docs_url if args.default_swagger_ui else None,
)
if args.cors:
app.add_middleware(
CORSMiddleware,
allow_origins=['*'],
allow_credentials=True,
allow_methods=['*'],
allow_headers=['*'],
)
logger.warning(
'CORS is enabled. This service is now accessible from any website!'
)
from jina.serve.runtimes.gateway.request_handling import RequestHandler
from jina.serve.stream import RequestStreamer
request_handler = RequestHandler(metrics_registry, args.name)
streamer = RequestStreamer(
args=args,
request_handler=request_handler.handle_request(
graph=topology_graph, connection_pool=connection_pool),
result_handler=request_handler.handle_result(),
)
streamer.Call = streamer.stream
@app.on_event('shutdown')
async def _shutdown():
await connection_pool.close()
openapi_tags = []
if not args.no_debug_endpoints:
openapi_tags.append({
'name':
'Debug',
'description':
'Debugging interface. In production, you should hide them by setting '
'`--no-debug-endpoints` in `Flow`/`Gateway`.',
})
from jina.serve.runtimes.gateway.http.models import JinaHealthModel
@app.get(
path='/',
summary='Get the health of Jina service',
response_model=JinaHealthModel,
)
async def _health():
"""
Get the health of this Jina service.
.. # noqa: DAR201
"""
return {}
@app.get(
path='/status',
summary='Get the status of Jina service',
response_model=JinaStatusModel,
tags=['Debug'],
)
async def _status():
"""
Get the status of this Jina service.
This is equivalent to running `jina -vf` from command line.
.. # noqa: DAR201
"""
_info = get_full_version()
return {
'jina': _info[0],
'envs': _info[1],
'used_memory': used_memory_readable(),
}
@app.post(
path='/post',
summary='Post a data request to some endpoint',
response_model=JinaResponseModel,
tags=['Debug']
# do not add response_model here, this debug endpoint should not restricts the response model
)
async def post(body: JinaEndpointRequestModel):
"""
Post a data request to some endpoint.
This is equivalent to the following:
from jina import Flow
f = Flow().add(...)
with f:
f.post(endpoint, ...)
.. # noqa: DAR201
.. # noqa: DAR101
"""
# The above comment is written in Markdown for better rendering in FastAPI
from jina.enums import DataInputType
bd = body.dict() # type: Dict
req_generator_input = bd
req_generator_input['data_type'] = DataInputType.DICT
if bd['data'] is not None and 'docs' in bd['data']:
req_generator_input['data'] = req_generator_input['data'][
'docs']
result = await _get_singleton_result(
request_generator(**req_generator_input))
return result
def expose_executor_endpoint(exec_endpoint, http_path=None, **kwargs):
"""Exposing an executor endpoint to http endpoint
:param exec_endpoint: the executor endpoint
:param http_path: the http endpoint
:param kwargs: kwargs accepted by FastAPI
"""
# set some default kwargs for richer semantics
# group flow exposed endpoints into `customized` group
kwargs['tags'] = kwargs.get('tags', ['Customized'])
kwargs['response_model'] = kwargs.get(
'response_model',
JinaResponseModel, # use standard response model by default
)
kwargs['methods'] = kwargs.get('methods', ['POST'])
@app.api_route(path=http_path or exec_endpoint,
name=http_path or exec_endpoint,
**kwargs)
async def foo(body: JinaRequestModel):
from jina.enums import DataInputType
bd = body.dict() if body else {'data': None}
bd['exec_endpoint'] = exec_endpoint
req_generator_input = bd
req_generator_input['data_type'] = DataInputType.DICT
if bd['data'] is not None and 'docs' in bd['data']:
req_generator_input['data'] = req_generator_input['data'][
'docs']
result = await _get_singleton_result(
request_generator(**req_generator_input))
return result
if not args.no_crud_endpoints:
openapi_tags.append({
'name':
'CRUD',
'description':
'CRUD interface. If your service does not implement those interfaces, you can should '
'hide them by setting `--no-crud-endpoints` in `Flow`/`Gateway`.',
})
crud = {
'/index': {
'methods': ['POST']
},
'/search': {
'methods': ['POST']
},
'/delete': {
'methods': ['DELETE']
},
'/update': {
'methods': ['PUT']
},
}
for k, v in crud.items():
v['tags'] = ['CRUD']
v['description'] = f'Post data requests to the Flow. Executors with `@requests(on="{k}")` will respond.'
expose_executor_endpoint(exec_endpoint=k, **v)
if openapi_tags:
app.openapi_tags = openapi_tags
if args.expose_endpoints:
endpoints = json.loads(args.expose_endpoints) # type: Dict[str, Dict]
for k, v in endpoints.items():
expose_executor_endpoint(exec_endpoint=k, **v)
if not args.default_swagger_ui:
async def _render_custom_swagger_html(req: Request) -> HTMLResponse:
import urllib.request
swagger_url = 'https://api.jina.ai/swagger'
req = urllib.request.Request(swagger_url,
headers={'User-Agent': 'Mozilla/5.0'})
with urllib.request.urlopen(req) as f:
return HTMLResponse(f.read().decode())
app.add_route(docs_url,
_render_custom_swagger_html,
include_in_schema=False)
if args.expose_graphql_endpoint:
with ImportExtensions(required=True):
from dataclasses import asdict
import strawberry
from docarray import DocumentArray
from docarray.document.strawberry_type import (
JSONScalar,
StrawberryDocument,
StrawberryDocumentInput,
)
from strawberry.fastapi import GraphQLRouter
async def get_docs_from_endpoint(data, target_executor, parameters,
exec_endpoint):
req_generator_input = {
'data': [asdict(d) for d in data],
'target_executor': target_executor,
'parameters': parameters,
'exec_endpoint': exec_endpoint,
'data_type': DataInputType.DICT,
}
if (req_generator_input['data'] is not None
and 'docs' in req_generator_input['data']):
req_generator_input['data'] = req_generator_input['data'][
'docs']
response = await _get_singleton_result(
request_generator(**req_generator_input))
return DocumentArray.from_dict(
response['data']).to_strawberry_type()
@strawberry.type
class Mutation:
@strawberry.mutation
async def docs(
self,
data: Optional[List[StrawberryDocumentInput]] = None,
target_executor: Optional[str] = None,
parameters: Optional[JSONScalar] = None,
exec_endpoint: str = '/search',
) -> List[StrawberryDocument]:
return await get_docs_from_endpoint(
data, target_executor, parameters, exec_endpoint)
@strawberry.type
class Query:
@strawberry.field
async def docs(
self,
data: Optional[List[StrawberryDocumentInput]] = None,
target_executor: Optional[str] = None,
parameters: Optional[JSONScalar] = None,
exec_endpoint: str = '/search',
) -> List[StrawberryDocument]:
return await get_docs_from_endpoint(
data, target_executor, parameters, exec_endpoint)
schema = strawberry.Schema(query=Query, mutation=Mutation)
app.include_router(GraphQLRouter(schema), prefix='/graphql')
async def _get_singleton_result(request_iterator) -> Dict:
"""
Streams results from AsyncPrefetchCall as a dict
:param request_iterator: request iterator, with length of 1
:return: the first result from the request iterator
"""
async for k in streamer.stream(request_iterator=request_iterator):
request_dict = k.to_dict()
return request_dict
return app
def get_fastapi_app(
args: 'argparse.Namespace',
topology_graph: 'TopologyGraph',
connection_pool: 'GrpcConnectionPool',
logger: 'JinaLogger',
metrics_registry: Optional['CollectorRegistry'] = None,
):
"""
Get the app from FastAPI as the REST interface.
:param args: passed arguments.
:param topology_graph: topology graph that manages the logic of sending to the proper executors.
:param connection_pool: Connection Pool to handle multiple replicas and sending to different of them
:param logger: Jina logger.
:param metrics_registry: optional metrics registry for prometheus used if we need to expose metrics from the executor or from the data request handler
:return: fastapi app
"""
with ImportExtensions(required=True):
from fastapi import FastAPI, Response, status
from fastapi.middleware.cors import CORSMiddleware
from jina.serve.runtimes.gateway.http.models import (
JinaEndpointRequestModel,
JinaRequestModel,
JinaResponseModel,
)
app = FastAPI(
title=args.title or 'My Jina Service',
description=args.description or
'This is my awesome service. You can set `title` and `description` in your `Flow` or `Gateway` '
'to customize the title and description.',
version=__version__,
)
if args.cors:
app.add_middleware(
CORSMiddleware,
allow_origins=['*'],
allow_credentials=True,
allow_methods=['*'],
allow_headers=['*'],
)
logger.warning(
'CORS is enabled. This service is accessible from any website!')
from jina.serve.runtimes.gateway.request_handling import RequestHandler
from jina.serve.stream import RequestStreamer
request_handler = RequestHandler(metrics_registry, args.name)
streamer = RequestStreamer(
args=args,
request_handler=request_handler.handle_request(
graph=topology_graph, connection_pool=connection_pool),
result_handler=request_handler.handle_result(),
)
streamer.Call = streamer.stream
@app.on_event('shutdown')
async def _shutdown():
await connection_pool.close()
openapi_tags = []
if not args.no_debug_endpoints:
openapi_tags.append({
'name':
'Debug',
'description':
'Debugging interface. In production, you should hide them by setting '
'`--no-debug-endpoints` in `Flow`/`Gateway`.',
})
from jina.serve.runtimes.gateway.http.models import JinaHealthModel
@app.get(
path='/',
summary='Get the health of Jina Gateway service',
response_model=JinaHealthModel,
)
async def _gateway_health():
"""
Get the health of this Gateway service.
.. # noqa: DAR201
"""
return {}
from docarray import DocumentArray
from jina.proto import jina_pb2
from jina.serve.executors import __dry_run_endpoint__
from jina.serve.runtimes.gateway.http.models import (
PROTO_TO_PYDANTIC_MODELS,
JinaInfoModel,
)
from jina.types.request.status import StatusMessage
@app.get(
path='/dry_run',
summary=
'Get the readiness of Jina Flow service, sends an empty DocumentArray to the complete Flow to '
'validate connectivity',
response_model=PROTO_TO_PYDANTIC_MODELS.StatusProto,
)
async def _flow_health():
"""
Get the health of the complete Flow service.
.. # noqa: DAR201
"""
da = DocumentArray()
try:
_ = await _get_singleton_result(
request_generator(
exec_endpoint=__dry_run_endpoint__,
data=da,
data_type=DataInputType.DOCUMENT,
))
status_message = StatusMessage()
status_message.set_code(jina_pb2.StatusProto.SUCCESS)
return status_message.to_dict()
except Exception as ex:
status_message = StatusMessage()
status_message.set_exception(ex)
return status_message.to_dict(use_integers_for_enums=True)
@app.get(
path='/status',
summary='Get the status of Jina service',
response_model=JinaInfoModel,
tags=['Debug'],
)
async def _status():
"""
Get the status of this Jina service.
This is equivalent to running `jina -vf` from command line.
.. # noqa: DAR201
"""
version, env_info = get_full_version()
for k, v in version.items():
version[k] = str(v)
for k, v in env_info.items():
env_info[k] = str(v)
return {'jina': version, 'envs': env_info}
@app.post(
path='/post',
summary='Post a data request to some endpoint',
response_model=JinaResponseModel,
tags=['Debug']
# do not add response_model here, this debug endpoint should not restricts the response model
)
async def post(
body: JinaEndpointRequestModel, response: Response
): # 'response' is a FastAPI response, not a Jina response
"""
Post a data request to some endpoint.
This is equivalent to the following:
from jina import Flow
f = Flow().add(...)
with f:
f.post(endpoint, ...)
.. # noqa: DAR201
.. # noqa: DAR101
"""
# The above comment is written in Markdown for better rendering in FastAPI
from jina.enums import DataInputType
bd = body.dict() # type: Dict
req_generator_input = bd
req_generator_input['data_type'] = DataInputType.DICT
if bd['data'] is not None and 'docs' in bd['data']:
req_generator_input['data'] = req_generator_input['data'][
'docs']
try:
result = await _get_singleton_result(
request_generator(**req_generator_input))
except InternalNetworkError as err:
import grpc
if err.code() == grpc.StatusCode.UNAVAILABLE:
response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
elif err.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
response.status_code = status.HTTP_504_GATEWAY_TIMEOUT
else:
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
result = bd # send back the request
result['header'] = _generate_exception_header(
err) # attach exception details to response header
logger.error(
f'Error while getting responses from deployments: {err.details()}'
)
return result
def _generate_exception_header(error: InternalNetworkError):
import traceback
from jina.proto.serializer import DataRequest
exception_dict = {
'name':
str(error.__class__),
'stacks': [
str(x)
for x in traceback.extract_tb(error.og_exception.__traceback__)
],
'executor':
'',
}
status_dict = {
'code': DataRequest().status.ERROR,
'description': error.details() if error.details() else '',
'exception': exception_dict,
}
header_dict = {'request_id': error.request_id, 'status': status_dict}
return header_dict
def expose_executor_endpoint(exec_endpoint, http_path=None, **kwargs):
"""Exposing an executor endpoint to http endpoint
:param exec_endpoint: the executor endpoint
:param http_path: the http endpoint
:param kwargs: kwargs accepted by FastAPI
"""
# set some default kwargs for richer semantics
# group flow exposed endpoints into `customized` group
kwargs['tags'] = kwargs.get('tags', ['Customized'])
kwargs['response_model'] = kwargs.get(
'response_model',
JinaResponseModel, # use standard response model by default
)
kwargs['methods'] = kwargs.get('methods', ['POST'])
@app.api_route(path=http_path or exec_endpoint,
name=http_path or exec_endpoint,
**kwargs)
async def foo(body: JinaRequestModel):
from jina.enums import DataInputType
bd = body.dict() if body else {'data': None}
bd['exec_endpoint'] = exec_endpoint
req_generator_input = bd
req_generator_input['data_type'] = DataInputType.DICT
if bd['data'] is not None and 'docs' in bd['data']:
req_generator_input['data'] = req_generator_input['data'][
'docs']
result = await _get_singleton_result(
request_generator(**req_generator_input))
return result
if not args.no_crud_endpoints:
openapi_tags.append({
'name':
'CRUD',
'description':
'CRUD interface. If your service does not implement those interfaces, you can should '
'hide them by setting `--no-crud-endpoints` in `Flow`/`Gateway`.',
})
crud = {
'/index': {
'methods': ['POST']
},
'/search': {
'methods': ['POST']
},
'/delete': {
'methods': ['DELETE']
},
'/update': {
'methods': ['PUT']
},
}
for k, v in crud.items():
v['tags'] = ['CRUD']
v['description'] = f'Post data requests to the Flow. Executors with `@requests(on="{k}")` will respond.'
expose_executor_endpoint(exec_endpoint=k, **v)
if openapi_tags:
app.openapi_tags = openapi_tags
if args.expose_endpoints:
endpoints = json.loads(args.expose_endpoints) # type: Dict[str, Dict]
for k, v in endpoints.items():
expose_executor_endpoint(exec_endpoint=k, **v)
if args.expose_graphql_endpoint:
with ImportExtensions(required=True):
from dataclasses import asdict
import strawberry
from docarray import DocumentArray
from docarray.document.strawberry_type import (
JSONScalar,
StrawberryDocument,
StrawberryDocumentInput,
)
from strawberry.fastapi import GraphQLRouter
async def get_docs_from_endpoint(data, target_executor, parameters,
exec_endpoint):
req_generator_input = {
'data': [asdict(d) for d in data],
'target_executor': target_executor,
'parameters': parameters,
'exec_endpoint': exec_endpoint,
'data_type': DataInputType.DICT,
}
if (req_generator_input['data'] is not None
and 'docs' in req_generator_input['data']):
req_generator_input['data'] = req_generator_input['data'][
'docs']
try:
response = await _get_singleton_result(
request_generator(**req_generator_input))
except InternalNetworkError as err:
logger.error(
f'Error while getting responses from deployments: {err.details()}'
)
raise err # will be handled by Strawberry
return DocumentArray.from_dict(
response['data']).to_strawberry_type()
@strawberry.type
class Mutation:
@strawberry.mutation
async def docs(
self,
data: Optional[List[StrawberryDocumentInput]] = None,
target_executor: Optional[str] = None,
parameters: Optional[JSONScalar] = None,
exec_endpoint: str = '/search',
) -> List[StrawberryDocument]:
return await get_docs_from_endpoint(
data, target_executor, parameters, exec_endpoint)
@strawberry.type
class Query:
@strawberry.field
async def docs(
self,
data: Optional[List[StrawberryDocumentInput]] = None,
target_executor: Optional[str] = None,
parameters: Optional[JSONScalar] = None,
exec_endpoint: str = '/search',
) -> List[StrawberryDocument]:
return await get_docs_from_endpoint(
data, target_executor, parameters, exec_endpoint)
schema = strawberry.Schema(query=Query, mutation=Mutation)
app.include_router(GraphQLRouter(schema), prefix='/graphql')
async def _get_singleton_result(request_iterator) -> Dict:
"""
Streams results from AsyncPrefetchCall as a dict
:param request_iterator: request iterator, with length of 1
:return: the first result from the request iterator
"""
async for k in streamer.stream(request_iterator=request_iterator):
request_dict = k.to_dict()
return request_dict
return app
def get_fastapi_app(
args: 'argparse.Namespace',
topology_graph: 'TopologyGraph',
connection_pool: 'GrpcConnectionPool',
logger: 'JinaLogger',
metrics_registry: Optional['CollectorRegistry'] = None,
):
"""
Get the app from FastAPI as the Websocket interface.
:param args: passed arguments.
:param topology_graph: topology graph that manages the logic of sending to the proper executors.
:param connection_pool: Connection Pool to handle multiple replicas and sending to different of them
:param logger: Jina logger.
:param metrics_registry: optional metrics registry for prometheus used if we need to expose metrics from the executor or from the data request handler
:return: fastapi app
"""
from jina.serve.runtimes.gateway.http.models import JinaEndpointRequestModel
with ImportExtensions(required=True):
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
class ConnectionManager:
def __init__(self):
self.active_connections: List[WebSocket] = []
self.protocol_dict: Dict[str, WebsocketSubProtocols] = {}
def get_client(self, websocket: WebSocket) -> str:
return f'{websocket.client.host}:{websocket.client.port}'
def get_subprotocol(self, headers: Dict):
try:
if 'sec-websocket-protocol' in headers:
subprotocol = WebsocketSubProtocols(
headers['sec-websocket-protocol'])
elif b'sec-websocket-protocol' in headers:
subprotocol = WebsocketSubProtocols(
headers[b'sec-websocket-protocol'].decode())
else:
subprotocol = WebsocketSubProtocols.JSON
logger.debug(
f'no protocol headers passed. Choosing default subprotocol {WebsocketSubProtocols.JSON}'
)
except Exception as e:
logger.debug(
f'got an exception while setting user\'s subprotocol, defaulting to JSON {e}'
)
subprotocol = WebsocketSubProtocols.JSON
return subprotocol
async def connect(self, websocket: WebSocket):
await websocket.accept()
subprotocol = self.get_subprotocol(dict(
websocket.scope['headers']))
logger.info(
f'client {websocket.client.host}:{websocket.client.port} connected '
f'with subprotocol {subprotocol}')
self.active_connections.append(websocket)
self.protocol_dict[self.get_client(websocket)] = subprotocol
def disconnect(self, websocket: WebSocket):
self.protocol_dict.pop(self.get_client(websocket))
self.active_connections.remove(websocket)
async def receive(self, websocket: WebSocket) -> Any:
subprotocol = self.protocol_dict[self.get_client(websocket)]
if subprotocol == WebsocketSubProtocols.JSON:
return await websocket.receive_json(mode='text')
elif subprotocol == WebsocketSubProtocols.BYTES:
return await websocket.receive_bytes()
async def iter(self, websocket: WebSocket) -> AsyncIterator[Any]:
try:
while True:
yield await self.receive(websocket)
except WebSocketDisconnect:
pass
async def send(self, websocket: WebSocket, data: DataRequest) -> None:
subprotocol = self.protocol_dict[self.get_client(websocket)]
if subprotocol == WebsocketSubProtocols.JSON:
return await websocket.send_json(data.to_dict(), mode='text')
elif subprotocol == WebsocketSubProtocols.BYTES:
return await websocket.send_bytes(data.to_bytes())
manager = ConnectionManager()
app = FastAPI()
from jina.serve.runtimes.gateway.request_handling import RequestHandler
from jina.serve.stream import RequestStreamer
request_handler = RequestHandler(metrics_registry, args.name)
streamer = RequestStreamer(
args=args,
request_handler=request_handler.handle_request(
graph=topology_graph, connection_pool=connection_pool),
result_handler=request_handler.handle_result(),
)
streamer.Call = streamer.stream
@app.on_event('shutdown')
async def _shutdown():
await connection_pool.close()
@app.websocket('/')
async def websocket_endpoint(websocket: WebSocket):
await manager.connect(websocket)
async def req_iter():
async for request in manager.iter(websocket):
if isinstance(request, dict):
if request == {}:
break
else:
# NOTE: Helps in converting camelCase to snake_case
req_generator_input = JinaEndpointRequestModel(
**request).dict()
req_generator_input['data_type'] = DataInputType.DICT
if request['data'] is not None and 'docs' in request[
'data']:
req_generator_input['data'] = req_generator_input[
'data']['docs']
# you can't do `yield from` inside an async function
for data_request in request_generator(
**req_generator_input):
yield data_request
elif isinstance(request, bytes):
if request == bytes(True):
break
else:
yield DataRequest(request)
try:
async for msg in streamer.stream(request_iterator=req_iter()):
await manager.send(websocket, msg)
except WebSocketDisconnect:
logger.info('Client successfully disconnected from server')
manager.disconnect(websocket)
return app
def get_fastapi_app(
args: 'argparse.Namespace',
topology_graph: 'TopologyGraph',
connection_pool: 'GrpcConnectionPool',
logger: 'JinaLogger',
):
"""
Get the app from FastAPI as the Websocket interface.
:param args: passed arguments.
:param topology_graph: topology graph that manages the logic of sending to the proper executors.
:param connection_pool: Connection Pool to handle multiple replicas and sending to different of them
:param logger: Jina logger.
:return: fastapi app
"""
with ImportExtensions(required=True):
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
class ConnectionManager:
def __init__(self):
self.active_connections: List[WebSocket] = []
async def connect(self, websocket: WebSocket):
await websocket.accept()
logger.debug(
f'client {websocket.client.host}:{websocket.client.port} connected'
)
self.active_connections.append(websocket)
def disconnect(self, websocket: WebSocket):
self.active_connections.remove(websocket)
manager = ConnectionManager()
app = FastAPI()
from jina.serve.stream import RequestStreamer
from jina.serve.runtimes.gateway.request_handling import (
handle_request,
handle_result,
)
streamer = RequestStreamer(
args=args,
request_handler=handle_request(graph=topology_graph,
connection_pool=connection_pool),
result_handler=handle_result,
)
streamer.Call = streamer.stream
@app.on_event('shutdown')
async def _shutdown():
await connection_pool.close()
@app.websocket('/')
async def websocket_endpoint(websocket: WebSocket):
await manager.connect(websocket)
async def req_iter():
async for request_bytes in websocket.iter_bytes():
if request_bytes == bytes(True):
break
yield DataRequest(request_bytes)
try:
async for msg in streamer.stream(request_iterator=req_iter()):
await websocket.send_bytes(bytes(msg))
except WebSocketDisconnect:
logger.debug('Client successfully disconnected from server')
manager.disconnect(websocket)
return app
def get_fastapi_app(
args: 'argparse.Namespace',
topology_graph: 'TopologyGraph',
connection_pool: 'GrpcConnectionPool',
logger: 'JinaLogger',
metrics_registry: Optional['CollectorRegistry'] = None,
):
"""
Get the app from FastAPI as the Websocket interface.
:param args: passed arguments.
:param topology_graph: topology graph that manages the logic of sending to the proper executors.
:param connection_pool: Connection Pool to handle multiple replicas and sending to different of them
:param logger: Jina logger.
:param metrics_registry: optional metrics registry for prometheus used if we need to expose metrics from the executor or from the data request handler
:return: fastapi app
"""
from jina.serve.runtimes.gateway.http.models import JinaEndpointRequestModel
with ImportExtensions(required=True):
from fastapi import FastAPI, Response, WebSocket, WebSocketDisconnect, status
class ConnectionManager:
def __init__(self):
self.active_connections: List[WebSocket] = []
self.protocol_dict: Dict[str, WebsocketSubProtocols] = {}
def get_client(self, websocket: WebSocket) -> str:
return f'{websocket.client.host}:{websocket.client.port}'
def get_subprotocol(self, headers: Dict):
try:
if 'sec-websocket-protocol' in headers:
subprotocol = WebsocketSubProtocols(
headers['sec-websocket-protocol']
)
elif b'sec-websocket-protocol' in headers:
subprotocol = WebsocketSubProtocols(
headers[b'sec-websocket-protocol'].decode()
)
else:
subprotocol = WebsocketSubProtocols.JSON
logger.debug(
f'no protocol headers passed. Choosing default subprotocol {WebsocketSubProtocols.JSON}'
)
except Exception as e:
logger.debug(
f'got an exception while setting user\'s subprotocol, defaulting to JSON {e}'
)
subprotocol = WebsocketSubProtocols.JSON
return subprotocol
async def connect(self, websocket: WebSocket):
await websocket.accept()
subprotocol = self.get_subprotocol(dict(websocket.scope['headers']))
logger.info(
f'client {websocket.client.host}:{websocket.client.port} connected '
f'with subprotocol {subprotocol}'
)
self.active_connections.append(websocket)
self.protocol_dict[self.get_client(websocket)] = subprotocol
def disconnect(self, websocket: WebSocket):
self.protocol_dict.pop(self.get_client(websocket))
self.active_connections.remove(websocket)
async def receive(self, websocket: WebSocket) -> Any:
subprotocol = self.protocol_dict[self.get_client(websocket)]
if subprotocol == WebsocketSubProtocols.JSON:
return await websocket.receive_json(mode='text')
elif subprotocol == WebsocketSubProtocols.BYTES:
return await websocket.receive_bytes()
async def iter(self, websocket: WebSocket) -> AsyncIterator[Any]:
try:
while True:
yield await self.receive(websocket)
except WebSocketDisconnect:
pass
async def send(
self, websocket: WebSocket, data: Union[DataRequest, StatusMessage]
) -> None:
subprotocol = self.protocol_dict[self.get_client(websocket)]
if subprotocol == WebsocketSubProtocols.JSON:
return await websocket.send_json(data.to_dict(), mode='text')
elif subprotocol == WebsocketSubProtocols.BYTES:
return await websocket.send_bytes(data.to_bytes())
manager = ConnectionManager()
app = FastAPI()
from jina.serve.runtimes.gateway.request_handling import RequestHandler
from jina.serve.stream import RequestStreamer
request_handler = RequestHandler(metrics_registry, args.name)
streamer = RequestStreamer(
args=args,
request_handler=request_handler.handle_request(
graph=topology_graph, connection_pool=connection_pool
),
result_handler=request_handler.handle_result(),
)
streamer.Call = streamer.stream
@app.get(
path='/',
summary='Get the health of Jina service',
)
async def _health():
"""
Get the health of this Jina service.
.. # noqa: DAR201
"""
return {}
@app.get(
path='/status',
summary='Get the status of Jina service',
)
async def _status():
"""
Get the status of this Jina service.
This is equivalent to running `jina -vf` from command line.
.. # noqa: DAR201
"""
version, env_info = get_full_version()
for k, v in version.items():
version[k] = str(v)
for k, v in env_info.items():
env_info[k] = str(v)
return {'jina': version, 'envs': env_info}
@app.on_event('shutdown')
async def _shutdown():
await connection_pool.close()
@app.websocket('/')
async def websocket_endpoint(
websocket: WebSocket, response: Response
): # 'response' is a FastAPI response, not a Jina response
await manager.connect(websocket)
async def req_iter():
async for request in manager.iter(websocket):
if isinstance(request, dict):
if request == {}:
break
else:
# NOTE: Helps in converting camelCase to snake_case
req_generator_input = JinaEndpointRequestModel(**request).dict()
req_generator_input['data_type'] = DataInputType.DICT
if request['data'] is not None and 'docs' in request['data']:
req_generator_input['data'] = req_generator_input['data'][
'docs'
]
# you can't do `yield from` inside an async function
for data_request in request_generator(**req_generator_input):
yield data_request
elif isinstance(request, bytes):
if request == bytes(True):
break
else:
yield DataRequest(request)
try:
async for msg in streamer.stream(request_iterator=req_iter()):
await manager.send(websocket, msg)
except InternalNetworkError as err:
import grpc
manager.disconnect(websocket)
fallback_msg = (
f'Connection to deployment at {err.dest_addr} timed out. You can adjust `timeout_send` attribute.'
if err.code() == grpc.StatusCode.DEADLINE_EXCEEDED
else f'Network error while connecting to deployment at {err.dest_addr}. It may be down.'
)
msg = (
err.details()
if _fits_ws_close_msg(
err.details()
) # some messages are too long for ws closing message
else fallback_msg
)
await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=msg)
except WebSocketDisconnect:
logger.info('Client successfully disconnected from server')
manager.disconnect(websocket)
async def _get_singleton_result(request_iterator) -> Dict:
"""
Streams results from AsyncPrefetchCall as a dict
:param request_iterator: request iterator, with length of 1
:return: the first result from the request iterator
"""
async for k in streamer.stream(request_iterator=request_iterator):
request_dict = k.to_dict()
return request_dict
from docarray import DocumentArray
from jina.proto import jina_pb2
from jina.serve.executors import __dry_run_endpoint__
from jina.serve.runtimes.gateway.http.models import PROTO_TO_PYDANTIC_MODELS
@app.get(
path='/dry_run',
summary='Get the readiness of Jina Flow service, sends an empty DocumentArray to the complete Flow to '
'validate connectivity',
response_model=PROTO_TO_PYDANTIC_MODELS.StatusProto,
)
async def _dry_run_http():
"""
Get the health of the complete Flow service.
.. # noqa: DAR201
"""
da = DocumentArray()
try:
_ = await _get_singleton_result(
request_generator(
exec_endpoint=__dry_run_endpoint__,
data=da,
data_type=DataInputType.DOCUMENT,
)
)
status_message = StatusMessage()
status_message.set_code(jina_pb2.StatusProto.SUCCESS)
return status_message.to_dict()
except Exception as ex:
status_message = StatusMessage()
status_message.set_exception(ex)
return status_message.to_dict(use_integers_for_enums=True)
@app.websocket('/dry_run')
async def websocket_endpoint(
websocket: WebSocket, response: Response
): # 'response' is a FastAPI response, not a Jina response
from jina.proto import jina_pb2
from jina.serve.executors import __dry_run_endpoint__
await manager.connect(websocket)
da = DocumentArray()
try:
async for _ in streamer.stream(
request_iterator=request_generator(
exec_endpoint=__dry_run_endpoint__,
data=da,
data_type=DataInputType.DOCUMENT,
)
):
pass
status_message = StatusMessage()
status_message.set_code(jina_pb2.StatusProto.SUCCESS)
await manager.send(websocket, status_message)
except InternalNetworkError as err:
manager.disconnect(websocket)
msg = (
err.details()
if _fits_ws_close_msg(err.details()) # some messages are too long
else f'Network error while connecting to deployment at {err.dest_addr}. It may be down.'
)
await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=msg)
except WebSocketDisconnect:
logger.info('Client successfully disconnected from server')
manager.disconnect(websocket)
except Exception as ex:
manager.disconnect(websocket)
status_message = StatusMessage()
status_message.set_exception(ex)
await manager.send(websocket, status_message)
return app