Repository: jina-ai/serve
Pull request: feat: add argument to control number of grpc retries
This allows the user to configure the number of retries that the gateway and head perform for each grpc call.
This includes the ability to disable retries altogether, which is useful when the user wants to user a service mesh level retry mechanism instead.

**ToDos:**

- [x] Add in Gateway
- [x] Add in Head
- [x] Test for gateway
- [x] Test for head
- [x] Add to docs

**Feature usage:**
Now the user can set the following:
```python
f = Flow()  # default policy: try every replica at least once, but perform at least 3 tries
f = Flow(retries=5)  # perform 5 retries, so 6 tries in total
f = Flow(retries=0)  # perform 0 retries, so 1 try in total
f = Flow(retries=-1)  # default policy again
```

Problem statement:
Allow user to disable retry mechanism
We have our own retry mechanism, which interferes with any service meshes retry mechanism (both would happen on top of each other, resulting in m*n retries, if we retry m times, and the service mesh n times).
This makes service mesh level retries not really usable.

We should allow the user to simply disable our own retry mechanism.

Hints from issue discussion:

Commit message: style: fix overload and cli autocomplete

Patch to review:
diff --git a/cli/autocomplete.py b/cli/autocomplete.py
index 19b9b8fd07cd6..cb0b104dfffe8 100644
--- a/cli/autocomplete.py
+++ b/cli/autocomplete.py
@@ -56,6 +56,7 @@
             '--port',
             '--monitoring',
             '--port-monitoring',
+            '--retries',
             '--install-requirements',
             '--force-update',
             '--force',
@@ -134,6 +135,7 @@
             '--port',
             '--monitoring',
             '--port-monitoring',
+            '--retries',
         ],
         'hub new': [
             '--help',
@@ -208,6 +210,7 @@
             '--port',
             '--monitoring',
             '--port-monitoring',
+            '--retries',
             '--install-requirements',
             '--force-update',
             '--force',
@@ -258,6 +261,7 @@
             '--port',
             '--monitoring',
             '--port-monitoring',
+            '--retries',
             '--install-requirements',
             '--force-update',
             '--force',
diff --git a/jina/orchestrate/flow/base.py b/jina/orchestrate/flow/base.py
index e4632339e4d0e..dae6d4e7b8a1e 100644
--- a/jina/orchestrate/flow/base.py
+++ b/jina/orchestrate/flow/base.py
@@ -166,6 +166,7 @@ def __init__(
         quiet: Optional[bool] = False,
         quiet_error: Optional[bool] = False,
         replicas: Optional[int] = 1,
+        retries: Optional[int] = None,
         runtime_cls: Optional[str] = 'GRPCGatewayRuntime',
         shards: Optional[int] = 1,
         ssl_certfile: Optional[str] = None,
@@ -242,6 +243,7 @@ def __init__(
         :param quiet: If set, then no log will be emitted from this object.
         :param quiet_error: If set, then exception stack information will not be added to the log
         :param replicas: The number of replicas in the deployment
+        :param retries: Number of retries per gRPC call. If None or <0, it defaults to max(3, num_replicas)
         :param runtime_cls: The runtime class to run inside the Pod
         :param shards: The number of shards in the deployment running at the same time. For more details check https://docs.jina.ai/fundamentals/flow/create-flow/#complex-flow-topologies
         :param ssl_certfile: the path to the certificate file
@@ -653,6 +655,7 @@ def add(
         quiet_error: Optional[bool] = False,
         quiet_remote_logs: Optional[bool] = False,
         replicas: Optional[int] = 1,
+        retries: Optional[int] = None,
         runtime_cls: Optional[str] = 'WorkerRuntime',
         shards: Optional[int] = 1,
         timeout_ctrl: Optional[int] = 60,
@@ -735,6 +738,7 @@ def add(
         :param quiet_error: If set, then exception stack information will not be added to the log
         :param quiet_remote_logs: Do not display the streaming of remote logs on local console
         :param replicas: The number of replicas in the deployment
+        :param retries: Number of retries per gRPC call. If None or <0, it defaults to max(3, num_replicas)
         :param runtime_cls: The runtime class to run inside the Pod
         :param shards: The number of shards in the deployment running at the same time. For more details check https://docs.jina.ai/fundamentals/flow/create-flow/#complex-flow-topologies
         :param timeout_ctrl: The timeout in milliseconds of the control request, -1 for waiting forever
diff --git a/jina/parsers/orchestrate/pod.py b/jina/parsers/orchestrate/pod.py
index 67f768e23d28b..538df7cae1c18 100644
--- a/jina/parsers/orchestrate/pod.py
+++ b/jina/parsers/orchestrate/pod.py
@@ -103,3 +103,11 @@ def mixin_pod_parser(parser):
         dest='port_monitoring',
         help=f'The port on which the prometheus server is exposed, default port is {__default_port_monitoring__} ',
     )
+
+    gp.add_argument(
+        '--retries',
+        type=int,
+        default=None,  # default prometheus server port
+        dest='retries',
+        help=f'Number of retries per gRPC call. If None or <0, it defaults to max(3, num_replicas)',
+    )
diff --git a/jina/serve/networking.py b/jina/serve/networking.py
index 4939302b6a8fb..4085af7236048 100644
--- a/jina/serve/networking.py
+++ b/jina/serve/networking.py
@@ -475,6 +475,7 @@ def send_request(
         polling_type: PollingType = PollingType.ANY,
         endpoint: Optional[str] = None,
         timeout: Optional[float] = None,
+        retries: Optional[int] = None,
     ) -> List[asyncio.Task]:
         """Send a single message to target via one or all of the pooled connections, depending on polling_type. Convenience function wrapper around send_request.
         :param request: a single request to send
@@ -484,6 +485,7 @@ def send_request(
         :param polling_type: defines if the message should be send to any or all pooled connections for the target
         :param endpoint: endpoint to target with the request
         :param timeout: timeout for sending the requests
+        :param retries: number of retries per gRPC call. If None or <0, it defaults to max(3, num_replicas)
         :return: list of asyncio.Task items for each send call
         """
         return self.send_requests(
@@ -494,6 +496,7 @@ def send_request(
             polling_type=polling_type,
             endpoint=endpoint,
             timeout=timeout,
+            retries=retries,
         )
 
     def send_requests(
@@ -505,6 +508,7 @@ def send_requests(
         polling_type: PollingType = PollingType.ANY,
         endpoint: Optional[str] = None,
         timeout: Optional[float] = None,
+        retries: Optional[int] = None,
     ) -> List[asyncio.Task]:
         """Send a request to target via one or all of the pooled connections, depending on polling_type
 
@@ -515,6 +519,7 @@ def send_requests(
         :param polling_type: defines if the request should be send to any or all pooled connections for the target
         :param endpoint: endpoint to target with the requests
         :param timeout: timeout for sending the requests
+        :param retries: number of retries per gRPC call. If None or <0, it defaults to max(3, num_replicas)
         :return: list of asyncio.Task items for each send call
         """
         results = []
@@ -532,7 +537,7 @@ def send_requests(
 
         for replica_list in connections:
             task = self._send_requests(
-                requests, replica_list, endpoint, timeout=timeout
+                requests, replica_list, endpoint, timeout=timeout, retries=retries
             )
             results.append(task)
 
@@ -568,6 +573,7 @@ def send_request_once(
         head: bool = False,
         shard_id: Optional[int] = None,
         timeout: Optional[float] = None,
+        retries: Optional[int] = None,
     ) -> asyncio.Task:
         """Send msg to target via only one of the pooled connections
         :param request: request to send
@@ -575,6 +581,7 @@ def send_request_once(
         :param head: If True it is send to the head, otherwise to the worker pods
         :param shard_id: Send to a specific shard of the deployment, ignored for polling ALL
         :param timeout: timeout for sending the requests
+        :param retries: number of retries per gRPC call. If None or <0, it defaults to max(3, num_replicas)
         :return: asyncio.Task representing the send call
         """
         return self.send_requests_once(
@@ -583,6 +590,7 @@ def send_request_once(
             head=head,
             shard_id=shard_id,
             timeout=timeout,
+            retries=retries,
         )
 
     def send_requests_once(
@@ -593,6 +601,7 @@ def send_requests_once(
         shard_id: Optional[int] = None,
         endpoint: Optional[str] = None,
         timeout: Optional[float] = None,
+        retries: Optional[int] = None,
     ) -> asyncio.Task:
         """Send a request to target via only one of the pooled connections
 
@@ -602,11 +611,14 @@ def send_requests_once(
         :param shard_id: Send to a specific shard of the deployment, ignored for polling ALL
         :param endpoint: endpoint to target with the requests
         :param timeout: timeout for sending the requests
+        :param retries: number of retries per gRPC call. If None or <0, it defaults to max(3, num_replicas)
         :return: asyncio.Task representing the send call
         """
         replicas = self._connections.get_replicas(deployment, head, shard_id)
         if replicas:
-            return self._send_requests(requests, replicas, endpoint, timeout=timeout)
+            return self._send_requests(
+                requests, replicas, endpoint, timeout=timeout, retries=retries
+            )
         else:
             self._logger.debug(
                 f'no available connections for deployment {deployment} and shard {shard_id}'
@@ -677,7 +689,7 @@ def _handle_aiorpcerror(
         retry_i: int = 0,
         request_id: str = '',
         dest_addr: Set[str] = {''},
-        num_retries: int = 3,
+        num_retries: int = 0,
     ):
         # connection failures and cancelled requests should be retried
         # all other cases should not be retried and will be raised immediately
@@ -690,7 +702,7 @@ def _handle_aiorpcerror(
             and e.code() != grpc.StatusCode.CANCELLED
         ):
             raise
-        elif e.code() == grpc.StatusCode.UNAVAILABLE and retry_i >= 2:
+        elif e.code() == grpc.StatusCode.UNAVAILABLE and retry_i >= num_retries - 1:
             self._logger.debug(f'GRPC call failed, retries exhausted')
             from jina.excepts import InternalNetworkError
 
@@ -712,13 +724,17 @@ def _send_requests(
         connections: ReplicaList,
         endpoint: Optional[str] = None,
         timeout: Optional[float] = None,
+        retries: Optional[int] = None,
     ) -> asyncio.Task:
         # this wraps the awaitable object from grpc as a coroutine so it can be used as a task
         # the grpc call function is not a coroutine but some _AioCall
         async def task_wrapper():
             metadata = (('endpoint', endpoint),) if endpoint else None
             tried_addresses = set()
-            num_retries = max(3, len(connections.get_all_connections()))
+            if retries is None or retries < 0:
+                num_retries = max(3, len(connections.get_all_connections()))
+            else:
+                num_retries = retries
             for i in range(num_retries):
                 current_connection = connections.get_next_connection()
                 tried_addresses.add(current_connection.address)
diff --git a/jina/serve/runtimes/gateway/__init__.py b/jina/serve/runtimes/gateway/__init__.py
index 89508f66d39b4..67f9e2cc13edd 100644
--- a/jina/serve/runtimes/gateway/__init__.py
+++ b/jina/serve/runtimes/gateway/__init__.py
@@ -43,6 +43,7 @@ def _set_topology_graph(self):
             graph_conditions,
             deployments_disable_reduce,
             timeout_send=self.timeout_send,
+            retries=self.args.retries,
         )
 
     def _set_connection_pool(self):
diff --git a/jina/serve/runtimes/gateway/graph/topology_graph.py b/jina/serve/runtimes/gateway/graph/topology_graph.py
index c1c4dc8e6b6cc..1434f69fa8650 100644
--- a/jina/serve/runtimes/gateway/graph/topology_graph.py
+++ b/jina/serve/runtimes/gateway/graph/topology_graph.py
@@ -36,6 +36,7 @@ def __init__(
             filter_condition: dict = None,
             reduce: bool = True,
             timeout_send: Optional[float] = None,
+            retries: Optional[int] = None,
         ):
             self.name = name
             self.outgoing_nodes = []
@@ -48,6 +49,7 @@ def __init__(
             self._filter_condition = filter_condition
             self._reduce = reduce
             self._timeout_send = timeout_send
+            self._retries = retries
 
         @property
         def leaf(self):
@@ -123,6 +125,7 @@ async def _wait_previous_and_send(
                             head=True,
                             endpoint=endpoint,
                             timeout=self._timeout_send,
+                            retries=self._retries,
                         )
                     except InternalNetworkError as err:
                         self._handle_internalnetworkerror(err)
@@ -237,6 +240,7 @@ def __init__(
         graph_conditions: Dict = {},
         deployments_disable_reduce: List[str] = [],
         timeout_send: Optional[float] = 1.0,
+        retries: Optional[int] = None,
         *args,
         **kwargs,
     ):
@@ -269,6 +273,7 @@ def __init__(
                 filter_condition=condition,
                 reduce=node_name not in deployments_disable_reduce,
                 timeout_send=timeout_send,
+                retries=retries,
             )
 
         for node_name, outgoing_node_names in graph_representation.items():


Code review comments:
diff_hunk: @@ -103,3 +103,11 @@ def mixin_pod_parser(parser):
         dest='port_monitoring',
         help=f'The port on which the prometheus server is exposed, default port is {__default_port_monitoring__} ',
     )
+
+    gp.add_argument(
+        '--retries',
+        type=int,
+        default=None,  # default prometheus server port
path: jina/parsers/orchestrate/pod.py
text: @user1:
it should be 3 right? plus wrong comment

@author:
Wrong comment, true, thanks for the catch.

But I would leave the default to `None`, which allows to keep the previous policy in place: Try at least three times, but if there are more than 3 replicas, try every replica at least once. That seems reasonable to me, because otherwise there might be a healthy replica but the call still fails. If the default is 3, we don't know if it was user set, and can't implement the default policy.

An alternative would be to always try every replica at least once, regardless of what the user puts in (`actual_num_retries = max(retries, num_replicas)`). But it feels a bit off to me to take that control away from the user.

Edit: `None` doesn't work, changed it to `-1` with the same semantics