def get_client(port):
args = set_client_cli_parser().parse_args([
'--host', 'localhost', '--port-expose',
str(port), '--return-results'
])
return Client(args)
def test_client_websocket(mocker, flow_with_rest_api_enabled):
with flow_with_rest_api_enabled:
time.sleep(0.5)
client = WebSocketClient(
set_client_cli_parser().parse_args(
[
'--host',
'localhost',
'--port-expose',
str(flow_with_rest_api_enabled.port_expose),
]
)
)
# Test that a regular index request triggers the correct callbacks
on_always_mock = mocker.Mock()
on_error_mock = mocker.Mock()
on_done_mock = mocker.Mock()
client.index(
iter([Document()]),
request_size=1,
on_always=on_always_mock,
on_error=on_error_mock,
on_done=on_done_mock,
)
on_always_mock.assert_called_once()
on_done_mock.assert_called_once()
on_error_mock.assert_not_called()
# Test that an empty index request does not trigger any callback and does not time out
mock = mocker.Mock()
client.index(
iter([()]), request_size=1, on_always=mock, on_error=mock, on_done=mock
)
mock.assert_not_called()
def _fetch_client(client: GatewayClients):
gateway_data_host = f'{os.getenv("JINA_GATEWAY_HOST") if os.getenv("JINA_GATEWAY_HOST") else "localhost"}'
gateway_data_port = f'{os.getenv("JINA_GATEWAY_PORT_EXPOSE") if os.getenv("JINA_GATEWAY_PORT_EXPOSE") else "23456"}'
args = set_client_cli_parser().parse_args(
['--host', gateway_data_host, '--port-expose',
str(gateway_data_port)])
return Client(args) if client == GatewayClients.GRPC else WebSocketClient(
args)
def main(task, load, nr, concurrency, req_size, dataset):
print(
f'task = {task}; load = {load}; nr = {nr}; concurrency = {concurrency}; \
req_size = {req_size}; dataset = {dataset}')
print(f'connecting to gRPC gateway at {FLOW_HOST}:{FLOW_PORT_GRPC}')
if task not in ['index', 'query']:
raise NotImplementedError(
f'unknown task: {task}. A valid task is either `index` or `query`.'
)
function = index
if task == 'query':
function = query
int_req_size = int(req_size)
grpc_args = set_client_cli_parser().parse_args(
['--host', FLOW_HOST, '--port-expose',
str(FLOW_PORT_GRPC)])
start_time = time.time()
time_end = start_time + load
print(f'Will end at {datetime.fromtimestamp(time_end).isoformat()}')
docs_generator = get_dataset_docs_gens(dataset)
if concurrency == 1:
print(f'Starting only one process. Not using multiprocessing...')
wrapper(grpc_args, docs_generator, 1, function, start_time, time_end,
int_req_size, dataset, nr)
else:
print(f'Using multiprocessing to start {concurrency} processes...')
processes = [
mp.Process(target=wrapper,
args=(grpc_args, docs_generator, id, function,
start_time, time_end, int_req_size, dataset, nr),
name=f'{function.__name__}-{id}')
for id in range(concurrency)
]
for p in processes:
p.start()
wait_secs = time_end - time.time()
if wait_secs > 0:
time.sleep(wait_secs)
for p in processes:
if p.is_alive():
print(f'Process {p.name} is still alive. Will wait...')
p.join()
# glob process logs file and sum the total processed
# by indexing and querying
end_time = time.time()
read_stats(time_end, op=task, time_spent=end_time - start_time)
async def _ping(host: str, port: int):
"""
Ping to check if we can connect to gateway via gRPC `host:port`
Note: Make sure Flow is running
"""
kwargs = {'port_expose': port, 'host': host}
_, args, _ = ArgNamespace.get_parsed_args(kwargs, set_client_cli_parser())
client = Client(args)
try:
# TODO: this introduces side-effect, need to be refactored. (2020.01.10)
client.index(input_fn=['abc'])
return {'status_code': status.HTTP_200_OK, 'detail': 'connected'}
except Exception:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f'Cannot connect to GRPC Server on {host}:{port}')
def parse_client(kwargs) -> Namespace:
"""
Parse the kwargs for the Client
:param kwargs: kwargs to be parsed
:return: parsed argument.
"""
kwargs = _parse_kwargs(kwargs)
args = ArgNamespace.kwargs2namespace(
kwargs, set_client_cli_parser(), warn_unknown=True
)
if not args.port:
args.port = (
__default_port_client__ if not args.tls else __default_port_tls_client__
)
return args
def __init__(
self,
args: Optional['argparse.Namespace'] = None,
**kwargs,
):
if args and isinstance(args, argparse.Namespace):
self.args = args
else:
self.args = ArgNamespace.kwargs2namespace(
kwargs, set_client_cli_parser(), warn_unknown=True
)
self.logger = JinaLogger(self.__class__.__name__, **vars(self.args))
if not self.args.proxy and os.name != 'nt':
# (Han 2020 12.12): gRPC channel is over HTTP2 and it does not work when we have proxy
# as many enterprise users are behind proxy, a quick way to
# surpass it is by temporally unset proxy. Please do NOT panic as it will NOT
# affect users os-level envs.
os.unsetenv('http_proxy')
os.unsetenv('https_proxy')
self._inputs = None
def main(task, port, load, nr, concurrency, req_size):
print(
f'task = {task}; port = {port}; load = {load}; nr = {nr}; concurrency = {concurrency}'
)
if task not in ['index', 'query']:
raise NotImplementedError(
f'unknown task: {task}. A valid task is either `index` or `query`.'
)
function = index
if task == 'query':
function = query
args = set_client_cli_parser().parse_args(
['--host', HOST, '--port-expose',
str(port)])
time_end = time.time() + load
print(f'Will end at {datetime.fromtimestamp(time_end).isoformat()}')
# needs to be a list otherwise it gets exhausted
docs = list(document_generator(nr))
# FIXME(cristianmtr): remote clients?
processes = [
mp.Process(target=wrapper,
args=(args, docs, id, function, time_end, req_size),
name=f'{function.__name__}-{id}')
for id in range(concurrency)
]
for p in processes:
p.start()
wait_secs = time_end - time.time()
if wait_secs > 0:
time.sleep(wait_secs)
for p in processes:
if p.is_alive():
print(f'Process {p.name} is still alive. Will wait...')
p.join()
def client():
args = set_client_cli_parser().parse_args(
['--host', 'localhost', '--port-expose', '45678'])
return Client(args)
def get_client(port):
args = set_client_cli_parser().parse_args(
['--host', 'localhost', '--port',
str(port)])
return Client(args)
def grpc_client():
args = set_client_cli_parser().parse_args(
['--host', 'localhost', '--port', '45678'])
return GRPCClient(args)
DenseNdArrayProto,
NdArrayProto,
SparseNdArrayProto,
NamedScoreProto,
DocumentProto,
RouteProto,
EnvelopeProto,
StatusProto,
MessageProto,
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()
PROTOBUF_TO_PYTHON_TYPE = {
FieldDescriptor.TYPE_INT32: int,
FieldDescriptor.TYPE_INT64: int,
FieldDescriptor.TYPE_UINT32: int,
FieldDescriptor.TYPE_UINT64: int,
FieldDescriptor.TYPE_SINT32: int,
FieldDescriptor.TYPE_SINT64: int,
FieldDescriptor.TYPE_BOOL: bool,
FieldDescriptor.TYPE_FLOAT: float,
FieldDescriptor.TYPE_DOUBLE: float,
FieldDescriptor.TYPE_FIXED32: float,
FieldDescriptor.TYPE_FIXED64: float,
FieldDescriptor.TYPE_SFIXED32: float,
FieldDescriptor.TYPE_SFIXED64: float,
dp = MessageToDict(pb_model(), including_default_value_fields=True)
all_fields = {
k: (name if k in ('chunks', 'matches') else type(v), Field(default=v))
for k, v in dp.items()
}
if pb_model == QueryLangProto:
all_fields['parameters'] = (Dict, Field(default={}))
return create_model(name, **all_fields)
JinaDocumentModel = build_model_from_pb('Document', DocumentProto)
JinaDocumentModel.update_forward_refs()
JinaQueryLangModel = build_model_from_pb('QueryLang', QueryLangProto)
default_request_size = set_client_cli_parser().parse_args([]).request_size
class JinaRequestModel(BaseModel):
"""
Jina request model.
The base model for Jina REST request.
"""
# To avoid an error while loading the request model schema on swagger, we've added an example.
data: Union[List[JinaDocumentModel], List[Dict[str, Any]], List[str], List[bytes]] = \
Field(..., example=[Document().dict()])
request_size: Optional[int] = default_request_size
mime_type: Optional[str] = ''
queryset: Optional[List[JinaQueryLangModel]] = None
data_type: DataInputType = DataInputType.AUTO
def parse_client(kwargs):
kwargs = _parse_kwargs(kwargs)
return ArgNamespace.kwargs2namespace(kwargs,
set_client_cli_parser(),
warn_unknown=True)
