Skip to content

Commit

Permalink
feat(rest): expose reload api to rest interface (#2301)
Browse files Browse the repository at this point in the history
  • Loading branch information
hanxiao authored Apr 11, 2021
1 parent c6ca3c8 commit d89564d
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 5 deletions.
21 changes: 21 additions & 0 deletions jina/peapods/runtimes/asyncio/rest/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ def get_fastapi_app(args: 'argparse.Namespace', logger: 'JinaLogger'):
JinaDeleteRequestModel,
JinaUpdateRequestModel,
JinaSearchRequestModel,
JinaReloadRequestModel,
)

app = FastAPI(
Expand Down Expand Up @@ -199,6 +200,26 @@ async def delete_api(body: JinaDeleteRequestModel):
result_in_stream(request_generator(**bd)), media_type='application/json'
)

@app.post(
path='/reload', summary='Reload the executor of certain Peas/Pods in the Flow'
)
async def reload_api(body: JinaReloadRequestModel):
"""
Reload the executor of certain peas/pods in the Flow
:param body: reload request.
:return: Response of the results.
"""
from .....clients import BaseClient

bd = body.dict()
bd['mode'] = RequestType.CONTROL
bd['command'] = 'RELOAD'
return StreamingResponse(
result_in_stream(request_generator(data=[], **bd)),
media_type='application/json',
)

async def result_in_stream(req_iter):
"""
Streams results from AsyncPrefetchCall as json
Expand Down
17 changes: 12 additions & 5 deletions jina/peapods/runtimes/asyncio/rest/models.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,12 @@
from enum import Enum
from collections import defaultdict
from datetime import datetime
from enum import Enum
from types import SimpleNamespace
from collections import defaultdict
from typing import Callable, Dict, Any, Optional, List, Union

from pydantic import Field, BaseModel, BaseConfig, create_model, root_validator
from google.protobuf.descriptor import Descriptor, FieldDescriptor
from google.protobuf.pyext.cpp_message import GeneratedProtocolMessageType

from jina.enums import DataInputType
from jina.types.document import Document
from jina.parsers import set_client_cli_parser
from jina.proto.jina_pb2 import (
DenseNdArrayProto,
Expand All @@ -24,6 +21,8 @@
RequestProto,
QueryLangProto,
)
from jina.types.document import Document
from pydantic import Field, BaseModel, BaseConfig, create_model, root_validator

DEFAULT_REQUEST_SIZE = set_client_cli_parser().parse_args([]).request_size
PROTO_TO_PYDANTIC_MODELS = SimpleNamespace()
Expand Down Expand Up @@ -238,6 +237,14 @@ class JinaStatusModel(BaseModel):
used_memory: str


class JinaReloadRequestModel(BaseModel):
"""
Jina control request model.
"""

targets: Union[str, List[str]]


class JinaRequestModel(BaseModel):
"""
Jina request model.
Expand Down
11 changes: 11 additions & 0 deletions tests/unit/flow/test_flow_reload.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,22 @@
import os

import pytest
import requests

from jina import Flow
from jina.executors import BaseExecutor


def test_flow_rest_reload():
f = Flow().add()
f.use_rest_gateway()
with f:
r = requests.post(
f'http://0.0.0.0:{f.port_expose}/reload', json={'targets': ['pod0']}
)
assert r.status_code == 200


@pytest.mark.skipif(
'GITHUB_WORKFLOW' in os.environ,
reason='skip the test on github as it will hang the whole CI, locally is fine',
Expand Down

0 comments on commit d89564d

Please sign in to comment.