Repository: jina-ai/serve
Pull request: feat: smart replica retry
Right now, when a replica fails (on local), the gateway will stubbornly retry to connect to that same replica, instead of trying its luck with one of the other replicas. This PR changes that, applying round robing scheduling also for retries, not just successful calls.

**ToDo:**

- [x] smart retry mechanism
- [x] test for retry mechanism
- [x] appropriate error message if all replicas fail
- [x] test for this failure case

Problem statement:
Smarter retry mechanism for Executors with replicas
Currently, the round robin scheduling between replicas does not apply when retrying to send a request, but only when a new request comes in. This means that if a replica is down, the gateway stubbornly tries to connect to the this same replica, ultimately failing the request despite other replicas being available.
This should be changed s.t. round robing scheduling is applied also to connection retries.

Hints from issue discussion:

Commit message: refactor: remove print

Patch to review:
diff --git a/jina/excepts.py b/jina/excepts.py
index 556dd858a63f6..18cc2c5aff5fc 100644
--- a/jina/excepts.py
+++ b/jina/excepts.py
@@ -1,4 +1,6 @@
 """This modules defines all kinds of exceptions raised in Jina."""
+from typing import Set, Union
+
 import grpc.aio
 
 
@@ -84,13 +86,13 @@ def __init__(
         self,
         og_exception: grpc.aio.AioRpcError,
         request_id: str = '',
-        dest_addr: str = '',
+        dest_addr: Union[str, Set[str]] = {''},
         details: str = '',
     ):
         """
         :param og_exception: the original exception that caused the network error
         :param request_id: id of the request that caused the error
-        :param dest_addr: destination (microservice) address of the problematic network call
+        :param dest_addr: destination (microservice) address(es) of the problematic network call(s)
         :param details: details of the error
         """
         self.og_exception = og_exception
diff --git a/jina/serve/networking.py b/jina/serve/networking.py
index c451a4df23284..813d030d4001a 100644
--- a/jina/serve/networking.py
+++ b/jina/serve/networking.py
@@ -3,7 +3,7 @@
 import ipaddress
 import os
 from collections import defaultdict
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Set, Tuple
 from urllib.parse import urlparse
 
 import grpc
@@ -310,6 +310,7 @@ def get_replicas(
             entity_id: Optional[int] = None,
             increase_access_count: bool = True,
         ) -> ReplicaList:
+            # returns all replicas of a given deployment, using a "random" shard
             if deployment in self._deployments:
                 type_ = 'heads' if head else 'shards'
                 if entity_id is None and head:
@@ -324,6 +325,8 @@ def get_replicas(
                 return None
 
         def get_replicas_all_shards(self, deployment: str) -> List[ReplicaList]:
+            # returns all replicas of a given deployment, for all available shards
+            # result is a list of 'shape' (num_shards, num_replicas), containing all replicas for all shards
             replicas = []
             if deployment in self._deployments:
                 for shard_id in self._deployments[deployment]['shards']:
@@ -527,18 +530,20 @@ def send_requests(
         results = []
         connections = []
         if polling_type == PollingType.ANY:
-            connection_list = self._connections.get_replicas(deployment, head, shard_id)
-            if connection_list:
-                connections.append(connection_list.get_next_connection())
+            replica_list = self._connections.get_replicas(deployment, head, shard_id)
+            if replica_list:
+                connections.append(replica_list)
         elif polling_type == PollingType.ALL:
-            connection_lists = self._connections.get_replicas_all_shards(deployment)
-            for connection_list in connection_lists:
-                connections.append(connection_list.get_next_connection())
+            shard_replica_lists = self._connections.get_replicas_all_shards(deployment)
+            for replica_list in shard_replica_lists:
+                connections.append(replica_list)
         else:
             raise ValueError(f'Unsupported polling type {polling_type}')
 
-        for connection in connections:
-            task = self._send_requests(requests, connection, endpoint, timeout=timeout)
+        for replica_list in connections:
+            task = self._send_requests(
+                requests, replica_list, endpoint, timeout=timeout
+            )
             results.append(task)
 
         return results
@@ -611,8 +616,7 @@ def send_requests_once(
         """
         replicas = self._connections.get_replicas(deployment, head, shard_id)
         if replicas:
-            connection = replicas.get_next_connection()
-            return self._send_requests(requests, connection, endpoint, timeout=timeout)
+            return self._send_requests(requests, replicas, endpoint, timeout=timeout)
         else:
             self._logger.debug(
                 f'no available connections for deployment {deployment} and shard {shard_id}'
@@ -682,7 +686,8 @@ def _handle_aiorpcerror(
         e: AioRpcError,
         retry_i: int = 0,
         request_id: str = '',
-        dest_addr: str = '',
+        dest_addr: Set[str] = {''},
+        num_retries: int = 3,
     ):
         # connection failures and cancelled requests should be retried
         # all other cases should not be retried and will be raised immediately
@@ -707,13 +712,14 @@ def _handle_aiorpcerror(
             )
         else:
             self._logger.debug(
-                f'GRPC call failed with code {e.code()}, retry attempt {retry_i + 1}/3'
+                f'GRPC call failed with code {e.code()}, retry attempt {retry_i + 1}/{num_retries}.'
+                f' Trying next replica, if available.'
             )
 
     def _send_requests(
         self,
         requests: List[Request],
-        connection: ConnectionStubs,
+        connections: ReplicaList,
         endpoint: Optional[str] = None,
         timeout: Optional[float] = None,
     ) -> asyncio.Task:
@@ -721,7 +727,11 @@ def _send_requests(
         # the grpc call function is not a coroutine but some _AioCall
         async def task_wrapper():
             metadata = (('endpoint', endpoint),) if endpoint else None
-            for i in range(3):
+            tried_addresses = set()
+            num_retries = max(3, len(connections.get_all_connections()))
+            for i in range(num_retries):
+                connection = connections.get_next_connection()
+                tried_addresses.add(connection.address)
                 try:
                     return await connection.send_requests(
                         requests=requests,
@@ -734,7 +744,8 @@ async def task_wrapper():
                         e=e,
                         retry_i=i,
                         request_id=requests[0].request_id,
-                        dest_addr=connection.address,
+                        dest_addr=tried_addresses,
+                        num_retries=num_retries,
                     )
 
         return asyncio.create_task(task_wrapper())
diff --git a/jina/serve/runtimes/gateway/graph/topology_graph.py b/jina/serve/runtimes/gateway/graph/topology_graph.py
index ffa166abe3deb..c1c4dc8e6b6cc 100644
--- a/jina/serve/runtimes/gateway/graph/topology_graph.py
+++ b/jina/serve/runtimes/gateway/graph/topology_graph.py
@@ -65,7 +65,7 @@ def _handle_internalnetworkerror(self, err):
             if err_code == grpc.StatusCode.UNAVAILABLE:
                 err._details = (
                     err.details()
-                    + f' |Gateway: Communication error with deployment {self.name} at address {err.dest_addr}. Head or worker may be down.'
+                    + f' |Gateway: Communication error with deployment {self.name} at address(es) {err.dest_addr}. Head or worker(s) may be down.'
                 )
                 raise err
             else:
diff --git a/jina/serve/runtimes/gateway/request_handling.py b/jina/serve/runtimes/gateway/request_handling.py
index 68568f5801094..78128df4284cc 100644
--- a/jina/serve/runtimes/gateway/request_handling.py
+++ b/jina/serve/runtimes/gateway/request_handling.py
@@ -89,7 +89,7 @@ def _get_all_nodes(node, accum, accum_names):
                 if err_code == grpc.StatusCode.UNAVAILABLE:
                     err._details = (
                         err.details()
-                        + f' |Gateway: Communication error with deployment at address {err.dest_addr}. Head or worker may be down.'
+                        + f' |Gateway: Communication error with deployment at address(es) {err.dest_addr}. Head or worker(s) may be down.'
                     )
                     raise err
                 else:
diff --git a/tests/integration/runtimes/test_network_failures.py b/tests/integration/runtimes/test_network_failures.py
index 04f6a78fb238a..664576c836355 100644
--- a/tests/integration/runtimes/test_network_failures.py
+++ b/tests/integration/runtimes/test_network_failures.py
@@ -71,11 +71,14 @@ def _send_request(gateway_port, protocol):
     )
 
 
-def _test_error(gateway_port, error_port, protocol):
+def _test_error(gateway_port, error_ports, protocol):
+    if not isinstance(error_ports, list):
+        error_ports = [error_ports]
     with pytest.raises(ConnectionError) as err_info:  # assert correct error is thrown
         _send_request(gateway_port, protocol)
-    # assert error message contains the port of the broken executor
-    assert str(error_port) in err_info.value.args[0]
+    # assert error message contains the port(s) of the broken executor(s)
+    for port in error_ports:
+        assert str(port) in err_info.value.args[0]
 
 
 @pytest.mark.parametrize(
@@ -141,7 +144,6 @@ async def test_runtimes_headless_topology(
             # so in this case, the actual request will fail, not the discovery, which is handled differently by Gateway
             worker_process.terminate()  # kill worker
             worker_process.join()
-            print(f'worker is alive: {worker_process.is_alive()}')
             assert not worker_process.is_alive()
         # ----------- 2. test that gateways remain alive -----------
         # just do the same again, expecting the same failure
@@ -411,3 +413,116 @@ async def test_runtimes_graphql(port_generator):
         worker_process.terminate()
         gateway_process.join()
         worker_process.join()
+
+
+@pytest.mark.asyncio
+async def test_replica_retry(port_generator):
+    # test that if one replica is down, the other replica(s) will be used
+    # create gateway and workers manually, then terminate worker process to provoke an error
+    worker_ports = [port_generator() for _ in range(3)]
+    worker0_port, worker1_port, worker2_port = worker_ports
+    gateway_port = port_generator()
+    graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
+    pod_addresses = f'{{"pod0": ["0.0.0.0:{worker0_port}", "0.0.0.0:{worker1_port}", "0.0.0.0:{worker2_port}"]}}'
+
+    worker_processes = []
+    for p in worker_ports:
+        worker_processes.append(_create_worker(p))
+        time.sleep(0.1)
+        AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
+            timeout=5.0,
+            ctrl_address=f'0.0.0.0:{p}',
+            ready_or_shutdown_event=multiprocessing.Event(),
+        )
+
+    gateway_process = _create_gateway(
+        gateway_port, graph_description, pod_addresses, 'grpc'
+    )
+    AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
+        timeout=5.0,
+        ctrl_address=f'0.0.0.0:{gateway_port}',
+        ready_or_shutdown_event=multiprocessing.Event(),
+    )
+
+    try:
+        # ----------- 1. ping Flow once to trigger endpoint discovery -----------
+        # we have to do this in a new process because otherwise grpc will be sad and everything will crash :(
+        p = multiprocessing.Process(target=_send_request, args=(gateway_port, 'grpc'))
+        p.start()
+        p.join()
+        assert p.exitcode == 0
+        # kill second worker, which would be responsible for the second call (round robin)
+        worker_processes[1].terminate()
+        worker_processes[1].join()
+        # ----------- 2. test that redundant replicas take over -----------
+        # we have to do this in a new process because otherwise grpc will be sad and everything will crash :(
+        p = multiprocessing.Process(target=_send_request, args=(gateway_port, 'grpc'))
+        p.start()
+        p.join()
+        assert p.exitcode == 0
+    except Exception:
+        assert False
+    finally:  # clean up runtimes
+        gateway_process.terminate()
+        gateway_process.join()
+        for p in worker_processes:
+            p.terminate()
+            p.join()
+
+
+@pytest.mark.asyncio
+async def test_replica_retry_all_fail(port_generator):
+    # test that if one replica is down, the other replica(s) will be used
+    # create gateway and workers manually, then terminate worker process to provoke an error
+    worker_ports = [port_generator() for _ in range(3)]
+    worker0_port, worker1_port, worker2_port = worker_ports
+    gateway_port = port_generator()
+    graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
+    pod_addresses = f'{{"pod0": ["0.0.0.0:{worker0_port}", "0.0.0.0:{worker1_port}", "0.0.0.0:{worker2_port}"]}}'
+
+    worker_processes = []
+    for p in worker_ports:
+        worker_processes.append(_create_worker(p))
+        time.sleep(0.1)
+        AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
+            timeout=5.0,
+            ctrl_address=f'0.0.0.0:{p}',
+            ready_or_shutdown_event=multiprocessing.Event(),
+        )
+
+    gateway_process = _create_gateway(
+        gateway_port, graph_description, pod_addresses, 'grpc'
+    )
+    AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
+        timeout=5.0,
+        ctrl_address=f'0.0.0.0:{gateway_port}',
+        ready_or_shutdown_event=multiprocessing.Event(),
+    )
+
+    try:
+        # ----------- 1. ping Flow once to trigger endpoint discovery -----------
+        # we have to do this in a new process because otherwise grpc will be sad and everything will crash :(
+        p = multiprocessing.Process(target=_send_request, args=(gateway_port, 'grpc'))
+        p.start()
+        p.join()
+        assert p.exitcode == 0
+        # kill all workers
+        for p in worker_processes:
+            p.terminate()
+            p.join()
+        # ----------- 2. test that call fails with informative error message -----------
+        # we have to do this in a new process because otherwise grpc will be sad and everything will crash :(
+        p = multiprocessing.Process(
+            target=_test_error, args=(gateway_port, worker_ports, 'grpc')
+        )
+        p.start()
+        p.join()
+        assert p.exitcode == 0
+    except Exception:
+        assert False
+    finally:  # clean up runtimes
+        gateway_process.terminate()
+        gateway_process.join()
+        for p in worker_processes:
+            p.terminate()
+            p.join()


Code review comments:
diff_hunk: @@ -310,6 +310,7 @@ def get_replicas(
             entity_id: Optional[int] = None,
             increase_access_count: bool = True,
         ) -> ReplicaList:
+            # returns all replicas of a given deployment, using a "random" shard
path: jina/serve/networking.py
text: @user1:
Not a `random` but a given one

diff_hunk: @@ -707,21 +712,26 @@ def _handle_aiorpcerror(
             )
         else:
             self._logger.debug(
-                f'GRPC call failed with code {e.code()}, retry attempt {retry_i + 1}/3'
+                f'GRPC call failed with code {e.code()}, retry attempt {retry_i + 1}/{num_retries}.'
+                f' Trying next replica, if available.'
             )
 
     def _send_requests(
         self,
         requests: List[Request],
-        connection: ConnectionStubs,
+        connections: ReplicaList,
         endpoint: Optional[str] = None,
         timeout: Optional[float] = 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
-            for i in range(3):
+            tried_addresses = set()
+            num_retries = max(3, len(connections.get_all_connections()))
+            for i in range(num_retries):
+                connection = connections.get_next_connection()
path: jina/serve/networking.py
text: @user1:
I find it really confusing to have `connection` for the single entity and `connections` for the group.
Maybe `current_connexion` ?