Repository: jina-ai/serve
Pull request: refactor: add disable_reduce args
# Refactor: replace has_used with disable_reduce argument

## Description

This PR adds the argument `disable-reduce`. This allows the user to disable the automatic reduce mechanism explicitly. By default it is false.
This only applies if there are no uses_before/uses_after defined. In case those are defined, they will be used for reducing the docs. This behavior is not changed.

Closes https://github.com/jina-ai/jina/issues/4381

Problem statement:
Change `_has_uses` logic in Head.
**Describe the bug**
This logic should be an independent `argument` saying `--reduce` that is set to True or False by the Head.

Also needs to check if both `Flow`,  `K8s` and `docker-compose` set properly this flag.

Hints from issue discussion:

Commit message: docs: add disable_reduce param

Patch to review:
diff --git a/cli/autocomplete.py b/cli/autocomplete.py
index 4338ccee257ae..11c37180188f5 100644
--- a/cli/autocomplete.py
+++ b/cli/autocomplete.py
@@ -97,6 +97,7 @@
             '--uses-before-address',
             '--uses-after-address',
             '--connection-list',
+            '--disable-reduce',
         ],
         'flow': [
             '--help',
@@ -172,6 +173,7 @@
             '--uses-before-address',
             '--uses-after-address',
             '--connection-list',
+            '--disable-reduce',
         ],
         'hub new': [
             '--help',
@@ -254,6 +256,7 @@
             '--uses-before-address',
             '--uses-after-address',
             '--connection-list',
+            '--disable-reduce',
         ],
         'deployment': [
             '--help',
@@ -305,6 +308,7 @@
             '--uses-before-address',
             '--uses-after-address',
             '--connection-list',
+            '--disable-reduce',
             '--uses-before',
             '--uses-after',
             '--external',
diff --git a/docs/fundamentals/flow/create-flow.md b/docs/fundamentals/flow/create-flow.md
index 73ac9bc366826..f36997553a0e8 100644
--- a/docs/fundamentals/flow/create-flow.md
+++ b/docs/fundamentals/flow/create-flow.md
@@ -425,6 +425,8 @@ This will get you the following output:
 
 So both `BarExecutor` and `BazExecutor` received only received a single `Document` from `FooExecutor` as they are run in parallel. The last Executor `executor3` will receive both DocumentArrays and merges them automatically.
 
+The automated merging can be disabled by setting `disable_reduce=True`. This can be useful when you need to provide your custom merge logic in a separate Executor. In this case the last `.add()` call would like `.add(needs=['barExecutor', 'bazExecutor'], uses=CustomMergeExecutor, disable_reduce=True)`
+
 ### Replicate Executors
 
 Replication can be used to create multiple copies of the same Executor. Each request in the Flow is then passed to only one replica (instance) of your Executor. This can be useful for a couple of challenges like performance and availability:
diff --git a/jina/helloworld/multimodal/flow-search.yml b/jina/helloworld/multimodal/flow-search.yml
index 7b3152e38483c..0d0edc68b4ed1 100644
--- a/jina/helloworld/multimodal/flow-search.yml
+++ b/jina/helloworld/multimodal/flow-search.yml
@@ -55,6 +55,7 @@ executors:
         py_modules:
           - ${{ ENV.PY_MODULE }}
     needs: [ textIndexer, imageIndexer ]
+    disable_reduce: True
   - name: keyvalueIndexer
     uses:
       jtype: KeyValueIndexer
diff --git a/jina/parsers/orchestrate/runtimes/head.py b/jina/parsers/orchestrate/runtimes/head.py
index d68b0e1a6c8d6..5ca88b6f07f45 100644
--- a/jina/parsers/orchestrate/runtimes/head.py
+++ b/jina/parsers/orchestrate/runtimes/head.py
@@ -1,4 +1,6 @@
-from jina.parsers.helper import add_arg_group
+import argparse
+
+from jina.parsers.helper import add_arg_group, _SHOW_ALL_ARGS
 
 
 def mixin_head_parser(parser):
@@ -25,3 +27,12 @@ def mixin_head_parser(parser):
         type=str,
         help='dictionary JSON with a list of connections to configure',
     )
+
+    gp.add_argument(
+        '--disable-reduce',
+        action='store_true',
+        default=False,
+        help='Disable the built-in reduce mechanism for merging parallel branches in a Flow'
+        if _SHOW_ALL_ARGS
+        else argparse.SUPPRESS,
+    )
diff --git a/jina/serve/networking.py b/jina/serve/networking.py
index 6f3047ad238d0..3bd914431d4b5 100644
--- a/jina/serve/networking.py
+++ b/jina/serve/networking.py
@@ -651,6 +651,45 @@ def send_request_sync(
                 if e.code() != grpc.StatusCode.UNAVAILABLE or i == 2:
                     raise
 
+    @staticmethod
+    def send_requests_sync(
+        requests: List[Request],
+        target: str,
+        timeout=100.0,
+        https=False,
+        root_certificates: Optional[str] = None,
+        endpoint: Optional[str] = None,
+    ) -> Request:
+        """
+        Sends a list of requests synchronically to the target via grpc
+
+        :param requests: the requests to send
+        :param target: where to send the request to, like 127.0.0.1:8080
+        :param timeout: timeout for the send
+        :param https: if True, use https for the grpc channel
+        :param root_certificates: the path to the root certificates for https, only used if https is True
+        :param endpoint: endpoint to target with the request
+
+        :returns: the response request
+        """
+
+        for i in range(3):
+            try:
+                with GrpcConnectionPool.get_grpc_channel(
+                    target,
+                    https=https,
+                    root_certificates=root_certificates,
+                ) as channel:
+                    metadata = (('endpoint', endpoint),) if endpoint else None
+                    stub = jina_pb2_grpc.JinaDataRequestRPCStub(channel)
+                    response, call = stub.process_data.with_call(
+                        requests, timeout=timeout, metadata=metadata
+                    )
+                    return response
+            except grpc.RpcError as e:
+                if e.code() != grpc.StatusCode.UNAVAILABLE or i == 2:
+                    raise
+
     @staticmethod
     def get_default_grpc_options():
         """
diff --git a/jina/serve/runtimes/head/__init__.py b/jina/serve/runtimes/head/__init__.py
index b06c980de4229..576b2e3ca7476 100644
--- a/jina/serve/runtimes/head/__init__.py
+++ b/jina/serve/runtimes/head/__init__.py
@@ -117,7 +117,7 @@ def __init__(
             self.connection_pool.add_connection(
                 deployment='uses_after', address=self.uses_after_address
             )
-        self._has_uses = args.uses is not None and args.uses != __default_executor__
+        self._reduce = not args.disable_reduce
 
     def _default_polling_dict(self, default_polling):
         return defaultdict(
@@ -252,7 +252,7 @@ async def _handle_data_request(
                 requests, deployment='uses_before'
             )
             requests = [response]
-        elif len(requests) > 1 and not self._has_uses:
+        elif len(requests) > 1 and self._reduce:
             requests = [DataRequestHandler.reduce_requests(requests)]
 
         worker_send_tasks = self.connection_pool.send_requests(
@@ -279,8 +279,14 @@ async def _handle_data_request(
             ) = await self.connection_pool.send_requests_once(
                 worker_results, deployment='uses_after'
             )
-        elif len(worker_results) > 1:
+        elif len(worker_results) > 1 and self._reduce:
             DataRequestHandler.reduce_requests(worker_results)
+        elif len(worker_results) > 1 and not self._reduce:
+            # worker returned multiple responsed, but the head is configured to skip reduction
+            # just concatenate the docs in this case
+            response_request.data.docs = DataRequestHandler.get_docs_from_request(
+                requests, field='docs'
+            )
 
         merged_metadata = self._merge_metadata(
             metadata, uses_after_metadata, uses_before_metadata
diff --git a/tests/integration/external_deployment/test_external_deployment.py b/tests/integration/external_deployment/test_external_deployment.py
index e64d80d68b8b3..fda7b86ff2d10 100644
--- a/tests/integration/external_deployment/test_external_deployment.py
+++ b/tests/integration/external_deployment/test_external_deployment.py
@@ -91,6 +91,7 @@ def test_flow_with_external_deployment(
 def test_two_flow_with_shared_external_deployment(
     external_deployment, external_deployment_args, input_docs, num_replicas, num_shards
 ):
+    external_deployment.head_args.disable_reduce = True
     with external_deployment:
         external_args = vars(external_deployment_args)
         del external_args['name']
@@ -295,6 +296,7 @@ def external_deployment_join_args(num_replicas, num_shards):
         str(num_replicas),
         '--polling',
         'all',
+        '--disable-reduce',
     ]
     return set_deployment_parser().parse_args(args)
 
diff --git a/tests/unit/orchestrate/flow/flow-orchestrate/test_flow_routing.py b/tests/unit/orchestrate/flow/flow-orchestrate/test_flow_routing.py
index 6c78667dc20a0..03579914839da 100644
--- a/tests/unit/orchestrate/flow/flow-orchestrate/test_flow_routing.py
+++ b/tests/unit/orchestrate/flow/flow-orchestrate/test_flow_routing.py
@@ -24,18 +24,24 @@ def add_text(self, docs, docs_matrix, **kwargs):
             docs[0].text = 'merged'
 
 
-def test_expected_messages_routing():
+@pytest.mark.parametrize('disable_reduce', [True, False])
+def test_expected_messages_routing(disable_reduce):
     f = (
         Flow()
         .add(name='foo', uses=SimplExecutor)
-        .add(name='bar', uses=MergeExecutor, needs=['foo', 'gateway'])
+        .add(
+            name='bar',
+            uses=MergeExecutor,
+            needs=['foo', 'gateway'],
+            disable_reduce=disable_reduce,
+        )
     )
 
     with f:
         docs = f.post(on='/index', inputs=[Document(text='1')])
         # there merge executor actually does not merge despite its name
-        assert len(docs) == 2
-        assert docs[0].text == 'merged'
+        assert len(docs) == 2 if disable_reduce else 1
+        assert docs[0].text == 'merged' if disable_reduce else '1'
 
 
 class SimpleAddExecutor(Executor):
@@ -58,7 +64,8 @@ def add_doc(self, docs, **kwargs):
         return docs
 
 
-def test_complex_flow():
+@pytest.mark.parametrize('disable_reduce', [True, False])
+def test_complex_flow(disable_reduce):
     f = (
         Flow()
         .add(name='first', uses=SimpleAddExecutor, needs=['gateway'])
@@ -75,12 +82,17 @@ def test_complex_flow():
             shards=3,
             needs=['second_shards_needs'],
         )
-        .add(name='merger', uses=MergeDocsExecutor, needs=['forth', 'third'])
+        .add(
+            name='merger',
+            uses=MergeDocsExecutor,
+            needs=['forth', 'third'],
+            disable_reduce=disable_reduce,
+        )
     )
 
     with f:
         docs = f.post(on='/index', inputs=[Document(text='1')])
-    assert len(docs) == 6
+    assert len(docs) == 6 if disable_reduce else 5
 
 
 class DynamicPollingExecutorDefaultNames(Executor):
diff --git a/tests/unit/serve/runtimes/head/test_head_runtime.py b/tests/unit/serve/runtimes/head/test_head_runtime.py
index 94f7d389057fa..0cf1a0911815f 100644
--- a/tests/unit/serve/runtimes/head/test_head_runtime.py
+++ b/tests/unit/serve/runtimes/head/test_head_runtime.py
@@ -1,7 +1,7 @@
 import asyncio
 import json
 import multiprocessing
-from copy import copy, deepcopy
+from copy import copy
 from multiprocessing import Process
 from typing import List
 
@@ -71,8 +71,12 @@ def test_control_message_processing():
     _destroy_runtime(args, cancel_event, runtime_thread)
 
 
-def test_message_merging():
-    args = set_pod_parser().parse_args([])
+@pytest.mark.parametrize('disable_reduce', [False, True])
+def test_message_merging(disable_reduce):
+    if not disable_reduce:
+        args = set_pod_parser().parse_args([])
+    else:
+        args = set_pod_parser().parse_args(['--disable-reduce'])
     args.polling = PollingType.ALL
     cancel_event, handle_queue, runtime_thread = _create_runtime(args)
 
@@ -82,12 +86,13 @@ def test_message_merging():
     _add_worker(args, 'ip3', shard_id=2)
     assert handle_queue.empty()
 
-    result = GrpcConnectionPool.send_request_sync(
-        _create_test_data_message(), f'{args.host}:{args.port_in}'
+    data_request = _create_test_data_message()
+    result = GrpcConnectionPool.send_requests_sync(
+        [data_request, data_request], f'{args.host}:{args.port_in}'
     )
     assert result
     assert _queue_length(handle_queue) == 3
-    assert len(result.response.docs) == 1
+    assert len(result.response.docs) == 2 if disable_reduce else 1
 
     _destroy_runtime(args, cancel_event, runtime_thread)
 


Code review comments:
diff_hunk: @@ -25,3 +27,12 @@ def mixin_head_parser(parser):
         type=str,
         help='dictionary JSON with a list of connections to configure',
     )
+
+    gp.add_argument(
+        '--disable-reduce',
+        action='store_true',
+        default=False,
+        help='Disable the built-in reduce mechanism for merging parallel branches in a Flow'
path: jina/parsers/orchestrate/runtimes/head.py
text: @user1:
I think this one should be shown and not hidden. 

I also would write something in the lines:

`help='Disable the built-in reduce mechanism, set this if the reduction is to be handled by the Executor connected to this Head'`