fix: uses=Executor should work with spawn start method (#3380)

jina/flow/base.py

@@ -712,6 +712,13 @@ def add(
             kwargs, parser, True, fallback_parsers=FALLBACK_PARSERS
         )
 
+        # Temporary workaround to re-import executor module when using mp spawn start
+        # method, so that the executor is re-registered with pyyaml. The re-import
+        # occurs because the class will be in the arguments passed to mp.Process.start
+        # method. A better solution probably involves deeper refactoring
+        if isinstance(kwargs.get('uses'), type(JAMLCompatible)):
+            args._exec_cls = kwargs['uses']
+
         if args.grpc_data_requests and args.runtime_cls == 'ZEDRuntime':
             args.runtime_cls = 'GRPCDataRuntime'
 

jina/jaml/__init__.py

@@ -507,6 +507,7 @@ def load_config(
         :param kwargs: kwargs for parse_config_source
         :return: :class:`JAMLCompatible` object
         """
+
         stream, s_path = parse_config_source(source, **kwargs)
         with stream as fp:
             # first load yml with no tag

jina/peapods/peas/__init__.py

@@ -52,6 +52,10 @@ def run(
     """
     logger = JinaLogger(name, **vars(args))
 
+    # Remove workaround used to re-import executor in spawn
+    if hasattr(args, '_exec_cls'):
+        del args._exec_cls
+
     def _unset_envs():
         if envs and args.runtime_backend != RuntimeBackendType.THREAD:
             for k in envs.keys():

tests/integration/multiprocessing_spawn/modules/exec.py

@@ -0,0 +1,5 @@
+from jina import Executor
+
+
+class Exec(Executor):
+    pass

tests/integration/multiprocessing_spawn/modules/main_cls.py

@@ -0,0 +1,15 @@
+from multiprocessing import get_start_method
+
+import jina
+
+
+def run():
+    from exec import Exec
+
+    with jina.Flow().add(uses=Exec) as f:
+        pass
+
+
+if __name__ == '__main__':
+    assert get_start_method() == 'spawn'
+    run()

tests/integration/multiprocessing_spawn/modules/main_empty.py

@@ -0,0 +1,15 @@
+from multiprocessing import get_start_method
+
+import jina
+
+
+def run():
+    from exec import Exec
+
+    with jina.Flow().add() as f:
+        pass
+
+
+if __name__ == '__main__':
+    assert get_start_method() == 'spawn'
+    run()

tests/integration/multiprocessing_spawn/modules/main_name.py

@@ -0,0 +1,15 @@
+from multiprocessing import get_start_method
+
+import jina
+
+
+def run():
+    from exec import Exec
+
+    with jina.Flow().add(uses="Exec") as f:
+        pass
+
+
+if __name__ == '__main__':
+    assert get_start_method() == 'spawn'
+    run()

tests/integration/multiprocessing_spawn/test_spawn.py

@@ -0,0 +1,41 @@
+"""Tests that flow can launch when using the spawn multiprocessing method
+
+Currently when using class name this still breaks (see xfail)
+"""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+
+import pytest
+
+
+def test_launch_spawn_empty():
+    subprocess.run(
+        [sys.executable, 'main_empty.py'],
+        check=True,
+        env={'JINA_MP_START_METHOD': 'spawn', 'PATH': os.environ['PATH']},
+        cwd=Path(__file__).parent / 'modules',
+    )
+
+
+def test_launch_spawn_cls():
+    print(sys.executable)
+    subprocess.run(
+        [sys.executable, 'main_cls.py'],
+        check=True,
+        env={'JINA_MP_START_METHOD': 'spawn', 'PATH': os.environ['PATH']},
+        cwd=Path(__file__).parent / 'modules',
+    )
+
+
+@pytest.mark.xfail(reason="re-importing not possible when given only exec name")
+def test_launch_spawn_name():
+    subprocess.run(
+        [sys.executable, 'main_name.py'],
+        check=True,
+        env={'JINA_MP_START_METHOD': 'spawn', 'PATH': os.environ['PATH']},
+        cwd=Path(__file__).parent / 'modules',
+    )