feat(flow): add install_requirments args when using jinahub (#3324)

Co-authored-by: Jina Dev Bot <dev-bot@jina.ai>

cli/autocomplete.py

@@ -320,6 +320,7 @@
             '--pea-id',
             '--pea-role',
             '--noblock-on-start',
+            '--install-requirements',
             '--uses-before',
             '--uses-after',
             '--parallel',

jina/flow/base.py

@@ -473,6 +473,7 @@ def add(
         host_in: Optional[str] = '0.0.0.0',
         host_out: Optional[str] = '0.0.0.0',
         hosts_in_connect: Optional[List[str]] = None,
+        install_requirements: Optional[bool] = False,
         log_config: Optional[str] = None,
         memory_hwm: Optional[int] = -1,
         name: Optional[str] = None,
@@ -539,6 +540,7 @@ def add(
         :param host_in: The host address for input, by default it is 0.0.0.0
         :param host_out: The host address for output, by default it is 0.0.0.0
         :param hosts_in_connect: The host address for input, by default it is 0.0.0.0
+        :param install_requirements: If set, install `requirements.txt` in the Hub Executor bundle
         :param log_config: The YAML config of the logger used in this object.
         :param memory_hwm: The memory high watermark of this pod in Gigabytes, pod will restart when this is reached. -1 means no restriction
         :param name: The name of this object.

jina/hubble/helper.py

@@ -10,9 +10,10 @@
 import tarfile
 import urllib
 import zipfile
+import pkg_resources
 from functools import lru_cache, wraps
 from pathlib import Path
-from typing import Tuple, Optional, Dict
+from typing import Tuple, Optional, Dict, List
 from urllib.parse import urlparse, urljoin
 from urllib.request import Request, urlopen
 
@@ -293,11 +294,33 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
-def install_requirements(requirements_file: 'Path'):
+def install_requirements(
+    requirements_file: 'Path', timeout: int = 1000, excludes: List[str] = ['jina']
+):
     """Install modules included in requirments file
-
     :param requirements_file: the requirements.txt file
+    :param timeout: the socket timeout (default = 1000s)
+    :param excludes: the excluded module dependencies
     """
+
+    with requirements_file.open() as requirements:
+        install_reqs = [
+            str(req)
+            for req in pkg_resources.parse_requirements(requirements)
+            if req.project_name not in excludes or len(req.extras) > 0
+        ]
+
+    if len(install_reqs) == 0:
+        return
+
     subprocess.check_call(
-        [sys.executable, '-m', 'pip', 'install', '-r', f'{requirements_file}']
+        [
+            sys.executable,
+            '-m',
+            'pip',
+            'install',
+            '--compile',
+            f'--default-timeout={timeout}',
+        ]
+        + install_reqs
     )

jina/parsers/peapods/pod.py

@@ -12,6 +12,12 @@ def mixin_base_pod_parser(parser):
     """
     gp = add_arg_group(parser, title='Pod')
 
+    gp.add_argument(
+        '--install-requirements',
+        action='store_true',
+        default=False,
+        help='If set, install `requirements.txt` in the Hub Executor bundle',
+    )
     gp.add_argument(
         '--uses-before',
         type=str,

jina/peapods/peas/helper.py

@@ -101,11 +101,14 @@ def update_runtime_cls(args, copy=False) -> 'Namespace':
     if _args.runtime_cls == 'ZEDRuntime' and _args.uses.startswith('docker://'):
         _args.runtime_cls = 'ContainerRuntime'
     if _args.runtime_cls == 'ZEDRuntime' and is_valid_huburi(_args.uses):
-        _args.uses = HubIO(
-            set_hub_pull_parser().parse_args([_args.uses, '--no-usage'])
-        ).pull()
+        _hub_args = [_args.uses, '--no-usage']
+        if _args.install_requirements:
+            _hub_args.append('--install-requirements')
+
+        _args.uses = HubIO(set_hub_pull_parser().parse_args(_hub_args)).pull()
         if _args.uses.startswith('docker://'):
             _args.runtime_cls = 'ContainerRuntime'
+
     if hasattr(_args, 'protocol'):
         _args.runtime_cls = gateway_runtime_dict[_args.protocol]
 

tests/integration/hub_usage/dummyhub/requirements.txt

@@ -0,0 +1 @@
+pytest

tests/integration/hub_usage/test_hub_usage.py

@@ -111,5 +111,5 @@ def _mock_fetch(name, tag=None, secret=None):
 
     monkeypatch.setattr(HubIO, 'fetch_meta', _mock_fetch)
 
-    with Flow().add(uses='jinahub://hello'):
+    with Flow().add(uses='jinahub://hello', install_requirements=True):
         pass

tests/unit/hubble/dummy_executor/requirements.txt

@@ -0,0 +1 @@
+pytest

tests/unit/hubble/test_helper.py

@@ -73,6 +73,12 @@ def test_unpack_package_unsupported(tmpdir):
         )
 
 
+def test_install_requirements():
+    helper.install_requirements(
+        Path(__file__).parent / 'dummy_executor' / 'requirements.txt'
+    )
+
+
 def test_disk_cache(tmpfile):
     raise_exception = True