Repository: jina-ai/serve
Pull request: feat: pass extra search paths to executor load config
add extra search paths to load_config in order to avoid relying just on the stack frame to retrieve executor yaml files
closes: https://github.com/jina-ai/jina/issues/3614

Problem statement:
Changes to adapt to FastAPI >= 0.70.0
Describe the feature
Due to failing tests/daemon/unit/api/endpoints/partial/test_flow.py::test_flow_api in CI, we found that the reason was the new release of FastAPI.

What we found is that the stack frame that we rely on to find the correct yml files to load did not mantain the tests... folder so I guess the frame is not mantained in a different process.

This is the symptom, and we limited the FastAPI version, however, the good thing to do is to rethink how all the paths are properly found, since relying on PATH env variable and the stack frame seems dangerous and unreliable.

Hints from issue discussion:

Commit message: test: add flow load config

Patch to review:
diff --git a/cli/autocomplete.py b/cli/autocomplete.py
index 956af4afb9345..3beb28de2539b 100644
--- a/cli/autocomplete.py
+++ b/cli/autocomplete.py
@@ -59,6 +59,7 @@
             '--static-routing-table',
             '--routing-table',
             '--dynamic-routing',
+            '--extra-search-paths',
             '--zmq-identity',
             '--port-ctrl',
             '--ctrl-with-ipc',
@@ -127,6 +128,7 @@
             '--static-routing-table',
             '--routing-table',
             '--dynamic-routing',
+            '--extra-search-paths',
             '--uses',
             '--env',
             '--inspect',
@@ -146,6 +148,7 @@
             '--static-routing-table',
             '--routing-table',
             '--dynamic-routing',
+            '--extra-search-paths',
             '--zmq-identity',
             '--port-ctrl',
             '--ctrl-with-ipc',
@@ -246,6 +249,7 @@
             '--static-routing-table',
             '--routing-table',
             '--dynamic-routing',
+            '--extra-search-paths',
             '--zmq-identity',
             '--port-ctrl',
             '--ctrl-with-ipc',
@@ -315,6 +319,7 @@
             '--static-routing-table',
             '--routing-table',
             '--dynamic-routing',
+            '--extra-search-paths',
             '--zmq-identity',
             '--port-ctrl',
             '--ctrl-with-ipc',
diff --git a/daemon/stores/partial.py b/daemon/stores/partial.py
index a55cc3bedccbb..2b66cc26c0745 100644
--- a/daemon/stores/partial.py
+++ b/daemon/stores/partial.py
@@ -106,7 +106,7 @@ def add(
             with open(args.uses) as yaml_file:
                 yaml_source = yaml_file.read()
 
-            self.object: Flow = Flow.load_config(yaml_source).build()
+            self.object: Flow = Flow.load_config(args.uses).build()
             self.object.workspace_id = jinad_args.workspace_id
             self.object.workspace = __partial_workspace__
             self.object.env = {'HOME': __partial_workspace__}
diff --git a/docs/fundamentals/flow/add-exec-to-flow.md b/docs/fundamentals/flow/add-exec-to-flow.md
index 384c5a01b7ccb..c462572c7dc2b 100644
--- a/docs/fundamentals/flow/add-exec-to-flow.md
+++ b/docs/fundamentals/flow/add-exec-to-flow.md
@@ -118,8 +118,97 @@ f = Flow().add(
 ```
 ````
 
+### Executor discovery
+By default, the Flow will attempt to retrieve the Executors' source files and yaml config files from the working 
+directory and other paths set in the `PATH` environment variable. If your Executor's source files and yaml config is 
+located elsewhere, you can specify the locations to be used in order to retrieve it using the parameter 
+`extra_search_paths`.
 
+For example, suppose we have the following project structure where `app/` represents the working directory:
+```
+.
+├── app
+│   └── ▶ main.py
+└── executor
+    ├── config.yml
+    └── my_executor.py
+```
+`executor/my_executor.py`:
+```python
+from jina import Executor, DocumentArray, requests
+
+class CustomExec(Executor):
+    @requests
+    def foo(self, docs: DocumentArray, **kwargs):
+        pass
+```
+
+`executor/config.yml`:
+```yaml
+jtype: CustomExec
+metas:
+  py_modules:
+    - executor.py
+```
+
+Now, in `app/main.py`, you can specify the directory of the Executor like so:
+```{code-block} python
+---
+emphasize-lines: 2
+---
+from jina import Flow, Document
+f = Flow(extra_search_paths=['../executor']).add(uses='config.yml')
+with f:
+    r = f.post('/', inputs=Document())
+```
 
+`extra_search_paths` will also be useful if you're specifying the Executor's config yaml directly in `main.py`:
+```{code-block} python
+---
+emphasize-lines: 2, 3, 4, 5, 6, 7
+---
+from jina import Flow, Document
+f = Flow(extra_search_paths=['../executor']).add(uses="""
+    jtype: CustomExec
+    metas:
+      py_modules:
+        - executor.py
+    """)
+with f:
+    r = f.post('/', inputs=Document())
+```
+
+````{admonition} Important
+:class: important
+If you are creating a Flow from a yaml config file which is located outside the working directory, you just need to 
+specify a correct relative or absolute path of the Flow's yaml config file and make all paths to Executor config files 
+relative to the Flow config file. The Flow will infer its config file location and add it to `extra_search_paths`:
+```
+.
+├── app
+│   └── ▶ main.py
+├── flow
+│   └── flow.yml
+└── executor
+    ├── config.yml
+    └── my_executor.py
+```
+`flow.yml`:
+```yaml
+jtype: Flow
+executors:
+  - name: executor
+    uses: ../executor/config.yml
+```
+
+`main.py`:
+```python
+from jina import Flow, Document
+f = Flow.load_config('../flow/flow.yml')
+with f:
+    r = f.post('/', inputs=Document())
+```
+````
 
 ## Add remote Executor
 
diff --git a/extra-requirements.txt b/extra-requirements.txt
index 272325e8c897b..85c0052f976b4 100644
--- a/extra-requirements.txt
+++ b/extra-requirements.txt
@@ -26,7 +26,7 @@
 # demo: required for the hello world demos
 
 scipy>=1.6.1:               devel
-fastapi<=0.68.2:            standard, daemon, devel, demo
+fastapi:            standard, daemon, devel, demo
 uvicorn[standard]>=0.14.0:  standard, daemon, devel, demo
 fluent-logger:              daemon, devel
 docker:                     standard, daemon, devel
diff --git a/jina/flow/base.py b/jina/flow/base.py
index e765c1b7758a0..1b4879a971a53 100644
--- a/jina/flow/base.py
+++ b/jina/flow/base.py
@@ -793,6 +793,7 @@ def add(
 
         args.k8s_namespace = self.args.name
         args.noblock_on_start = True
+        args.extra_search_paths = self.args.extra_search_paths
 
         # BACKWARDS COMPATIBILITY:
         # We assume that this is used in a search Flow if replicas and shards are used
diff --git a/jina/jaml/__init__.py b/jina/jaml/__init__.py
index 18e54186cfe2f..d63ad6a13850d 100644
--- a/jina/jaml/__init__.py
+++ b/jina/jaml/__init__.py
@@ -1,3 +1,4 @@
+import inspect
 import sys
 import os
 import re
@@ -480,6 +481,7 @@ def load_config(
         override_with: Optional[Dict] = None,
         override_metas: Optional[Dict] = None,
         override_requests: Optional[Dict] = None,
+        extra_search_paths: Optional[List[str]] = None,
         **kwargs,
     ) -> 'JAMLCompatible':
         """A high-level interface for loading configuration with features
@@ -529,11 +531,16 @@ def load_config(
         :param override_with: dictionary of parameters to overwrite from the default config's with field
         :param override_metas: dictionary of parameters to overwrite from the default config's metas field
         :param override_requests: dictionary of parameters to overwrite from the default config's requests field
+        :param extra_search_paths: extra paths used when looking for executor yaml files
         :param kwargs: kwargs for parse_config_source
         :return: :class:`JAMLCompatible` object
         """
+        if isinstance(source, str) and os.path.exists(source):
+            extra_search_paths = (extra_search_paths or []) + [os.path.dirname(source)]
 
-        stream, s_path = parse_config_source(source, **kwargs)
+        stream, s_path = parse_config_source(
+            source, extra_search_paths=extra_search_paths, **kwargs
+        )
         with stream as fp:
             # first load yml with no tag
             no_tag_yml = JAML.load_no_tags(fp)
@@ -570,15 +577,26 @@ def _delitem(
             if substitute:
                 # expand variables
                 no_tag_yml = JAML.expand_dict(no_tag_yml, context)
+
             if allow_py_modules:
+                _extra_search_paths = extra_search_paths or []
                 load_py_modules(
                     no_tag_yml,
-                    extra_search_paths=(os.path.dirname(s_path),) if s_path else None,
+                    extra_search_paths=(_extra_search_paths + [os.path.dirname(s_path)])
+                    if s_path
+                    else _extra_search_paths,
                 )
 
             from ..flow.base import Flow
 
-            if issubclass(cls, Flow):
+            if issubclass(cls, Flow) and no_tag_yml.get('jtype') != 'Flow {}':
+                # only needed for Flow
+                if no_tag_yml.get('with') is None:
+                    no_tag_yml['with'] = {}
+                no_tag_yml['with']['extra_search_paths'] = (
+                    no_tag_yml['with'].get('extra_search_paths') or []
+                ) + (extra_search_paths or [])
+
                 tag_yml = JAML.unescape(
                     JAML.dump(no_tag_yml),
                     include_unknown_tags=False,
diff --git a/jina/jaml/helper.py b/jina/jaml/helper.py
index 5aeb775f9a41e..279321083f4cb 100644
--- a/jina/jaml/helper.py
+++ b/jina/jaml/helper.py
@@ -2,7 +2,7 @@
 import json
 import os
 import warnings
-from typing import Union, TextIO, Dict, Tuple, Optional
+from typing import Union, TextIO, Dict, Tuple, Optional, List
 
 from yaml import MappingNode
 from yaml.composer import Composer
@@ -113,6 +113,7 @@ def parse_config_source(
     allow_class_type: bool = True,
     allow_dict: bool = True,
     allow_json: bool = True,
+    extra_search_paths: Optional[List[str]] = None,
     *args,
     **kwargs,
 ) -> Tuple[TextIO, Optional[str]]:
@@ -127,6 +128,7 @@ def parse_config_source(
     :param allow_class_type: flag
     :param allow_dict: flag
     :param allow_json: flag
+    :param extra_search_paths: extra paths to search for
     :param args: unused
     :param kwargs: unused
     :return: a tuple, the first element is the text stream, the second element is the file path associate to it
@@ -145,7 +147,7 @@ def parse_config_source(
         # already a readable stream
         return path, None
     elif allow_yaml_file and is_yaml_filepath(path):
-        comp_path = complete_path(path)
+        comp_path = complete_path(path, extra_search_paths)
         return open(comp_path, encoding='utf8'), comp_path
     elif allow_raw_yaml_content and path.lstrip().startswith(('!', 'jtype')):
         # possible YAML content
@@ -170,7 +172,7 @@ def parse_config_source(
         )
 
 
-def complete_path(path: str, extra_search_paths: Optional[Tuple[str]] = None) -> str:
+def complete_path(path: str, extra_search_paths: Optional[List[str]] = None) -> str:
     """
     Complete the path of file via searching in abs and relative paths.
 
@@ -188,7 +190,7 @@ def complete_path(path: str, extra_search_paths: Optional[Tuple[str]] = None) ->
         raise FileNotFoundError(f'can not find {path}')
 
 
-def _search_file_in_paths(path, extra_search_paths: Optional[Tuple[str]] = None):
+def _search_file_in_paths(path, extra_search_paths: Optional[List[str]] = None):
     """
     Search in all dirs of the PATH environment variable and all dirs of files used in the call stack.
 
@@ -217,7 +219,7 @@ def _search_file_in_paths(path, extra_search_paths: Optional[Tuple[str]] = None)
             return _p
 
 
-def load_py_modules(d: Dict, extra_search_paths: Optional[Tuple[str]] = None) -> None:
+def load_py_modules(d: Dict, extra_search_paths: Optional[List[str]] = None) -> None:
     """
     Find 'py_modules' in the dict recursively and then load them.
 
diff --git a/jina/jaml/parsers/flow/v1.py b/jina/jaml/parsers/flow/v1.py
index 1eb9de305f21c..1d8aa97042c7f 100644
--- a/jina/jaml/parsers/flow/v1.py
+++ b/jina/jaml/parsers/flow/v1.py
@@ -57,6 +57,7 @@ def parse(self, cls: type, data: Dict) -> 'Flow':
         :return: the Flow YAML parser given the syntax version number
         """
         p = data.get('with', {})  # type: Dict[str, Any]
+
         a = p.pop('args') if 'args' in p else ()
         k = p.pop('kwargs') if 'kwargs' in p else {}
         # maybe there are some hanging kwargs in "parameters"
diff --git a/jina/parsers/peapods/base.py b/jina/parsers/peapods/base.py
index b7dd48e2dafdc..c2824d484748c 100644
--- a/jina/parsers/peapods/base.py
+++ b/jina/parsers/peapods/base.py
@@ -111,3 +111,13 @@ def mixin_base_ppr_parser(parser, with_identity: bool = True):
         if _SHOW_ALL_ARGS
         else argparse.SUPPRESS,
     )
+
+    parser.add_argument(
+        '--extra-search-paths',
+        type=str,
+        default=[],
+        nargs='*',
+        help='Extra search paths to be used when loading modules.'
+        if _SHOW_ALL_ARGS
+        else argparse.SUPPRESS,
+    )
diff --git a/jina/peapods/runtimes/request_handlers/data_request_handler.py b/jina/peapods/runtimes/request_handlers/data_request_handler.py
index 315b429a4faee..212f7f005666b 100644
--- a/jina/peapods/runtimes/request_handlers/data_request_handler.py
+++ b/jina/peapods/runtimes/request_handlers/data_request_handler.py
@@ -75,6 +75,7 @@ def _load_executor(self):
                 override_metas=self.args.uses_metas,
                 override_requests=self.args.uses_requests,
                 runtime_args=vars(self.args),
+                extra_search_paths=self.args.extra_search_paths,
             )
         except BadConfigSource as ex:
             self.logger.error(
diff --git a/jina/resources/extra-requirements.txt b/jina/resources/extra-requirements.txt
index 272325e8c897b..85c0052f976b4 100644
--- a/jina/resources/extra-requirements.txt
+++ b/jina/resources/extra-requirements.txt
@@ -26,7 +26,7 @@
 # demo: required for the hello world demos
 
 scipy>=1.6.1:               devel
-fastapi<=0.68.2:            standard, daemon, devel, demo
+fastapi:            standard, daemon, devel, demo
 uvicorn[standard]>=0.14.0:  standard, daemon, devel, demo
 fluent-logger:              daemon, devel
 docker:                     standard, daemon, devel
diff --git a/tests/unit/flow-construct/executor/config.yml b/tests/unit/flow-construct/executor/config.yml
new file mode 100644
index 0000000000000..53b0cc6a9a24b
--- /dev/null
+++ b/tests/unit/flow-construct/executor/config.yml
@@ -0,0 +1,4 @@
+jtype: CustomExec
+metas:
+  py_modules:
+    - executor.py
diff --git a/tests/unit/flow-construct/executor/executor.py b/tests/unit/flow-construct/executor/executor.py
new file mode 100644
index 0000000000000..3732f6c4e500d
--- /dev/null
+++ b/tests/unit/flow-construct/executor/executor.py
@@ -0,0 +1,8 @@
+from jina import Executor, DocumentArray, requests
+
+
+class CustomExec(Executor):
+    @requests
+    def foo(self, docs: DocumentArray, **kwargs):
+        for doc in docs:
+            doc.text = 'done'
diff --git a/tests/unit/flow-construct/flow/flow.yml b/tests/unit/flow-construct/flow/flow.yml
new file mode 100644
index 0000000000000..0529345b260bf
--- /dev/null
+++ b/tests/unit/flow-construct/flow/flow.yml
@@ -0,0 +1,4 @@
+jtype: Flow
+executors:
+  - name: executor
+    uses: ../executor/config.yml
diff --git a/tests/unit/flow-construct/test_flow.py b/tests/unit/flow-construct/test_flow.py
index 6358768f00728..3480accb231a0 100644
--- a/tests/unit/flow-construct/test_flow.py
+++ b/tests/unit/flow-construct/test_flow.py
@@ -879,3 +879,33 @@ def foo(self, **kwargs):
         assert r[0].parameters['a'] == 1.0
         r = f.post('/', parameters={}, return_results=True)
         assert r[0].parameters['a'] == 1.0
+
+
+def test_flow_load_executor_extra_search_paths():
+    f = Flow(extra_search_paths=[os.path.join(cur_dir, 'executor')]).add(
+        uses="""
+    jtype: CustomExec
+    metas:
+      py_modules:
+        - executor.py
+    """
+    )
+    with f:
+        r = f.post('/', inputs=Document(), return_results=True)
+    assert r[0].docs[0].text == 'done'
+
+
+def test_flow_load_executor_yaml_extra_search_paths():
+    f = Flow(extra_search_paths=[os.path.join(cur_dir, 'executor')]).add(
+        uses='config.yml'
+    )
+    with f:
+        r = f.post('/', inputs=Document(), return_results=True)
+    assert r[0].docs[0].text == 'done'
+
+
+def test_flow_load_yaml_extra_search_paths():
+    f = Flow.load_config(os.path.join(cur_dir, 'flow/flow.yml'))
+    with f:
+        r = f.post('/', inputs=Document(), return_results=True)
+    assert r[0].docs[0].text == 'done'


Code review comments:
diff_hunk: @@ -1,3 +1,4 @@
+import inspect
path: jina/jaml/__init__.py
text: @user1:
remove useless import