Repository: jina-ai/serve
Pull request: fix(daemon): remote flow with parallel>1

Problem statement:
`JinaD` doesn't set `workspace` properly
While starting a Flow, `JinaD` sets Executor workspace to `<jinad-root-workspace>/<workspace-id>` (e.g.- `/tmp/jinad/jworkspace-0764ff64-ba12-437c-99f0-4a34f135180f` which is only accessible inside dockerhost and not on localhost. Other Flows on the same workspace cannot use the already indexed files. We should set the workspace to `/workspace` as it is mounted from localhost.

Hints from issue discussion:

Commit message: Merge branch 'master' into fix-remote-flow-replicas

Patch to review:
diff --git a/daemon/clients/flows.py b/daemon/clients/flows.py
index a7c0d9219e46e..67cf31b4cc22b 100644
--- a/daemon/clients/flows.py
+++ b/daemon/clients/flows.py
@@ -67,6 +67,9 @@ async def create(
                     f'{self._kind.title()} creation failed as: {error_msg}'
                 )
                 return error_msg
+            self._logger.success(
+                f'Remote Flow created successfully with id {response_json}'
+            )
             return response_json
 
     @if_alive
diff --git a/daemon/helper.py b/daemon/helper.py
index 5ee1358be269c..d914297376e11 100644
--- a/daemon/helper.py
+++ b/daemon/helper.py
@@ -99,8 +99,8 @@ def error_msg_from(response: Dict) -> str:
     :param response: dict response
     :return: prettified response string
     """
-    assert 'detail' in response, '\'detail\' not found in response'
-    assert 'body' in response, '\'body\' not found in response'
+    if 'detail' not in response or 'body' not in response:
+        return response
     if response['detail'] == PartialDaemon400Exception.__name__:
         return response['body']
     return (
diff --git a/daemon/stores/partial.py b/daemon/stores/partial.py
index 745bef973d4c2..b1255e8d7591f 100644
--- a/daemon/stores/partial.py
+++ b/daemon/stores/partial.py
@@ -8,7 +8,7 @@
 from jina import Flow, __docker_host__
 from jina.logging.logger import JinaLogger
 
-from .. import jinad_args
+from .. import jinad_args, __partial_workspace__
 from ..models import GATEWAY_RUNTIME_DICT
 from ..models.enums import UpdateOperation
 from ..models.ports import Ports, PortMappings
@@ -108,6 +108,7 @@ def add(
 
             self.object: Flow = Flow.load_config(yaml_source).build()
             self.object.workspace_id = jinad_args.workspace_id
+            self.object.workspace = __partial_workspace__
 
             for pod in self.object._pod_nodes.values():
                 runtime_cls = update_runtime_cls(pod.args, copy=True).runtime_cls
@@ -144,6 +145,12 @@ def add(
                             pod.replicas_args = CompoundPod._set_replica_args(
                                 pod.args, pod.head_args, pod.tail_args
                             )
+                else:
+                    # avoid setting runs_in_docker for Pods with parallel > 1 and using `ZEDRuntime`
+                    # else, replica-peas would try connecting to head/tail-pea via __docker_host__
+                    if runtime_cls == 'ZEDRuntime' and pod.args.parallel > 1:
+                        pod.args.runs_in_docker = False
+                        pod.update_pea_args()
 
             self.object = self.object.__enter__()
         except Exception as e:
diff --git a/jina/flow/base.py b/jina/flow/base.py
index 14c8c3b7a599a..877983aaceacf 100644
--- a/jina/flow/base.py
+++ b/jina/flow/base.py
@@ -1584,6 +1584,12 @@ def workspace(self) -> str:
         .. # noqa: DAR201"""
         return os.path.abspath(self.args.workspace or './')
 
+    @workspace.setter
+    def workspace(self, value: str):
+        self.args.workspace = value
+        for k, p in self:
+            p.args.workspace = value
+
     @property
     def workspace_id(self) -> Dict[str, str]:
         """Get all Pods' ``workspace_id`` values in a dict
diff --git a/tests/distributed/test_topologies/test_topologies.py b/tests/distributed/test_topologies/test_topologies.py
index 0dc7bf7f8658c..6026de9b0f8dc 100644
--- a/tests/distributed/test_topologies/test_topologies.py
+++ b/tests/distributed/test_topologies/test_topologies.py
@@ -1,8 +1,11 @@
 import os
 
 import pytest
+import numpy as np
 
-from jina import Flow, Document
+from daemon.clients import JinaDClient
+from daemon import __partial_workspace__
+from jina import Flow, Document, Client, __default_host__
 
 cur_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -36,6 +39,7 @@ def test_r_l_simple(silent_log, parallels, mocker):
         f.index(
             inputs=(Document(text='hello') for _ in range(NUM_DOCS)),
             on_done=response_mock,
+            show_progress=True,
         )
 
     response_mock.assert_called()
@@ -50,6 +54,7 @@ def test_l_r_simple(parallels, mocker):
         f.index(
             inputs=(Document(text='hello') for _ in range(NUM_DOCS)),
             on_done=response_mock,
+            show_progress=True,
         )
     response_mock.assert_called()
 
@@ -68,6 +73,7 @@ def test_r_l_r_simple(parallels, mocker):
         f.index(
             inputs=(Document(text='hello') for _ in range(NUM_DOCS)),
             on_done=response_mock,
+            show_progress=True,
         )
     response_mock.assert_called()
 
@@ -86,6 +92,7 @@ def test_r_r_r_simple(parallels, mocker):
         f.index(
             inputs=(Document(text='hello') for _ in range(NUM_DOCS)),
             on_done=response_mock,
+            show_progress=True,
         )
     response_mock.assert_called()
 
@@ -99,6 +106,7 @@ def test_l_r_l_simple(parallels, mocker):
         f.index(
             inputs=(Document(text='hello') for _ in range(NUM_DOCS)),
             on_done=response_mock,
+            show_progress=True,
         )
     response_mock.assert_called()
 
@@ -144,3 +152,56 @@ def test_complex_needs(parallels, mocker):
             on_done=response_mock,
         )
     response_mock.assert_called()
+
+
+@pytest.mark.parametrize('parallel', [1, 2])
+def test_remote_flow_local_executors(mocker, parallel):
+    response_mock = mocker.Mock()
+    client = JinaDClient(host=__default_host__, port=8000)
+    workspace_id = client.workspaces.create(paths=[os.path.join(cur_dir, 'yamls')])
+
+    GATEWAY_LOCAL_GATEWAY = 'flow_glg.yml'
+    GATEWAY_LOCAL_LOCAL_GATEWAY = 'flow_gllg.yml'
+
+    for flow_yaml in [
+        GATEWAY_LOCAL_GATEWAY,
+        GATEWAY_LOCAL_LOCAL_GATEWAY,
+    ]:
+        flow_id = client.flows.create(
+            workspace_id=workspace_id, filename=flow_yaml, envs={'PARALLEL': parallel}
+        )
+        args = client.flows.get(flow_id)['arguments']['object']['arguments']
+        Client(
+            host=__default_host__,
+            port=args['port_expose'],
+            protocol=args['protocol'],
+        ).post(
+            on='/',
+            inputs=(Document(blob=np.random.random([1, 100])) for _ in range(NUM_DOCS)),
+            on_done=response_mock,
+            show_progress=True,
+        )
+        response_mock.assert_called()
+        assert client.flows.delete(flow_id)
+
+    assert client.workspaces.delete(workspace_id)
+
+
+def test_remote_workspace_value():
+    HOST = __default_host__  # '3.208.18.63'
+    client = JinaDClient(host=HOST, port=8000)
+    workspace_id = client.workspaces.create(paths=[os.path.join(cur_dir, 'yamls')])
+    flow_id = client.flows.create(
+        workspace_id=workspace_id, filename='flow_workspace_validate.yml'
+    )
+    args = client.flows.get(flow_id)['arguments']['object']['arguments']
+    response = Client(
+        host=HOST, port=args['port_expose'], protocol=args['protocol']
+    ).post(on='/', inputs=[Document()], show_progress=True, return_results=True)
+    assert (
+        response[0]
+        .data.docs[0]
+        .text.startswith(f'{__partial_workspace__}/WorkspaceValidator/0')
+    )
+    assert client.flows.delete(flow_id)
+    assert client.workspaces.delete(workspace_id)
diff --git a/tests/distributed/test_topologies/yamls/executor.py b/tests/distributed/test_topologies/yamls/executor.py
new file mode 100644
index 0000000000000..4d45b9d82013c
--- /dev/null
+++ b/tests/distributed/test_topologies/yamls/executor.py
@@ -0,0 +1,7 @@
+from jina import Executor, requests, DocumentArray
+
+
+class WorkspaceValidator(Executor):
+    @requests
+    def foo(self, docs: DocumentArray, *args, **kwargs):
+        docs[0].text = self.workspace
diff --git a/tests/distributed/test_topologies/yamls/flow_glg.yml b/tests/distributed/test_topologies/yamls/flow_glg.yml
new file mode 100644
index 0000000000000..2b334f9ce17fe
--- /dev/null
+++ b/tests/distributed/test_topologies/yamls/flow_glg.yml
@@ -0,0 +1,8 @@
+jtype: Flow
+version: 1
+with:
+  port_expose: 9000
+  protocol: http
+executors:
+  - name: local
+    parallel: $PARALLEL
diff --git a/tests/distributed/test_topologies/yamls/flow_gllg.yml b/tests/distributed/test_topologies/yamls/flow_gllg.yml
new file mode 100644
index 0000000000000..468ddd626b56e
--- /dev/null
+++ b/tests/distributed/test_topologies/yamls/flow_gllg.yml
@@ -0,0 +1,10 @@
+jtype: Flow
+version: 1
+with:
+  port_expose: 9000
+  protocol: http
+executors:
+  - name: local1
+    parallel: $PARALLEL
+  - name: local2
+    parallel: $PARALLEL
diff --git a/tests/distributed/test_topologies/yamls/flow_workspace_validate.yml b/tests/distributed/test_topologies/yamls/flow_workspace_validate.yml
new file mode 100644
index 0000000000000..48ab9d33d2cf6
--- /dev/null
+++ b/tests/distributed/test_topologies/yamls/flow_workspace_validate.yml
@@ -0,0 +1,9 @@
+jtype: Flow
+version: 1
+with:
+  port_expose: 9000
+  protocol: http
+executors:
+  - name: local
+    uses: WorkspaceValidator
+    py_module: executor.py


Code review comments:
diff_hunk: @@ -144,6 +145,12 @@ def add(
                             pod.replicas_args = CompoundPod._set_replica_args(
                                 pod.args, pod.head_args, pod.tail_args
                             )
+                else:
+                    # avoid setting runs_in_docker for Pods with parallel > 1 and using `ZEDRuntime`
+                    # else, replica-peas would try connecting to head/tail-pea via __docker_host__
+                    if runtime_cls == 'ZEDRuntime' and pod.args.parallel > 1:
path: daemon/stores/partial.py
text: @user1:
can u add this to `GRPCDataRuntime` if possible?