feat: parse arbitrary type docker kwargs

daemon/models/custom.py

@@ -55,13 +55,13 @@ def get_pydantic_fields(config: Union[dict, argparse.ArgumentParser]):
                                   description=arg['help'])
             all_options[arg_key] = (arg_type, current_field)
 
-    # Issue: For all args that have a default value which comes from a function call (get_random_identity() or random_port()),
-    # 1st pydantic model invocation sets these default values, means build_pydantic_model(...) sets the args, not SinglePodModel()
-    # In case of multiple Pods, port conflict happens because of same port set as default in both.
-    # TODO(Deepankar): Add support for `default_factory` for default args that are functions
+    # Issue: For all args that have a default value which comes from a function call (get_random_identity() or
+    # random_port()), 1st pydantic model invocation sets these default values, means build_pydantic_model(...) sets
+    # the args, not SinglePodModel() In case of multiple Pods, port conflict happens because of same port set as
+    # default in both. TODO(Deepankar): Add support for `default_factory` for default args that are functions
     if isinstance(config, argparse.ArgumentParser):
         # Ignoring first 3 as they're generic args
-        from jina.parsers.helper import KVAppendAction
+        from jina.parsers.helper import KVAppendAction, DockerKwargsAppendAction
         for arg in config._actions[3:]:
             arg_key = arg.dest
             arg_type = arg.type
@@ -72,7 +72,7 @@ def get_pydantic_fields(config: Union[dict, argparse.ArgumentParser]):
             if hasattr(arg_type, '__self__'):
                 arg_type = type(arg.default) if arg.default else int
             arg_type = str if isinstance(arg_type, argparse.FileType) else arg_type
-            arg_type = dict if type(arg) == KVAppendAction else arg_type
+            arg_type = dict if (type(arg) == KVAppendAction or type(arg) == DockerKwargsAppendAction) else arg_type
 
             current_field = Field(default=arg.default,
                                   example=arg.default,

jina/flow/base.py

@@ -228,7 +228,6 @@ def add(self,
             num_part=len(needs)
         ))
 
-
         parser = set_pod_parser()
         if pod_role == PodRoleType.GATEWAY:
             parser = set_gateway_parser()

jina/parsers/helper.py

@@ -30,6 +30,30 @@ def __call__(self, parser, args, values, option_string=None):
         setattr(args, self.dest, d)
 
 
+class DockerKwargsAppendAction(argparse.Action):
+    """
+    argparse action to split an argument into KEY: VALUE form
+    on the first : and append to a dictionary.
+    This is used for setting up arbitrary kwargs for docker sdk
+    """
+
+    def __call__(self, parser, args, values, option_string=None):
+        import json
+        d = getattr(args, self.dest) or {}
+
+        for value in values:
+            try:
+                d.update(json.loads(value))
+            except json.JSONDecodeError:
+                try:
+                    (k, v) = value.split(':', 1)
+                except ValueError:
+                    raise argparse.ArgumentTypeError(f'could not parse argument \"{values[0]}\" as k:v format')
+                # transform from text to actual type (int, list, etc...)
+                d[k] = json.loads(v)
+        setattr(args, self.dest, d)
+
+
 class _ColoredHelpFormatter(argparse.ArgumentDefaultsHelpFormatter):
     class _Section(object):
 

jina/parsers/peapods/runtimes/container.py

@@ -1,4 +1,4 @@
-from ...helper import add_arg_group, KVAppendAction
+from ...helper import add_arg_group, DockerKwargsAppendAction
 
 
 def mixin_container_runtime_parser(parser):
@@ -11,8 +11,8 @@ def mixin_container_runtime_parser(parser):
     gp.add_argument('--entrypoint', type=str,
                     help='the entrypoint command overrides the ENTRYPOINT in docker image. '
                          'when not set then the docker image ENTRYPOINT takes effective.')
-    gp.add_argument('--docker-kwargs', action=KVAppendAction,
-                    metavar='KEY=VALUE', nargs='*',
+    gp.add_argument('--docker-kwargs', action=DockerKwargsAppendAction,
+                    metavar='KEY:VALUE', nargs='*',
                     help='dictionary of kwargs arguments that will be passed to docker sdk when starting the docker '
                          'image. Some of the arguments can be found in docker sdk documentation ('
                          'https://docker-py.readthedocs.io/en/stable/)')

tests/unit/parsers/peapods/runtimes/test_runtime_parser.py

@@ -13,6 +13,11 @@ def test_runtime_parser():
     args = parser.parse_args([])
     assert args.docker_kwargs is None
 
-    args = parser.parse_args(['--docker-kwargs', 'hello=0', 'bye=1'])
+    args = parser.parse_args(['--docker-kwargs', 'hello: 0', 'bye: 1'])
+    assert args.docker_kwargs == {'hello': 0, 'bye': 1}
+
+    args = parser.parse_args(['--docker-kwargs', 'hello: "0"', 'bye: "1"'])
     assert args.docker_kwargs == {'hello': '0', 'bye': '1'}
 
+    args = parser.parse_args(['--docker-kwargs', 'environment: ["VAR1=BAR", "VAR2=FOO"]', 'hello: 0'])
+    assert args.docker_kwargs == {'environment': ['VAR1=BAR', 'VAR2=FOO'], 'hello': 0}

tests/unit/peapods/runtimes/container/flow.yml

@@ -0,0 +1,8 @@
+!Flow
+pods:
+  pod1:
+    show-exc-info: true
+    uses: docker://jinahub/pod
+    docker-kwargs:
+      hello: 0
+      environment: ['VAR1=BAR', 'VAR2=FOO']

tests/unit/peapods/runtimes/container/test_container_runtime.py

@@ -180,12 +180,14 @@ def __init__(self):
             pass
 
         def run(self, *args, **kwargs):
-            mock_kwargs = {k: kwargs[k] for k in ['hello', 'ports']}
+            mock_kwargs = {k: kwargs[k] for k in ['hello', 'ports', 'environment']}
             mock(**mock_kwargs)
             assert 'ports' in kwargs
             assert kwargs['ports'] is None
+            assert 'environment' in kwargs
+            assert kwargs['environment'] == ['VAR1=BAR', 'VAR2=FOO']
             assert 'hello' in kwargs
-            assert kwargs['hello'] == '0'
+            assert kwargs['hello'] == 0
 
     class MockClient:
 
@@ -205,8 +207,13 @@ def containers(self):
 
     monkeypatch.setattr(docker, 'from_env', MockClient)
     args = set_pea_parser().parse_args(
-        ['--uses', 'docker://jinahub/pod', '--docker-kwargs', 'hello=0'])
+        ['--uses', 'docker://jinahub/pod', '--docker-kwargs', 'hello: 0', 'environment: ["VAR1=BAR", "VAR2=FOO"]'])
     runtime = ContainerRuntime(args)
     runtime._docker_run(replay=False)
-    expected_args = {'hello': '0', 'ports': None}
+    expected_args = {'hello': 0, 'ports': None, 'environment': ['VAR1=BAR', 'VAR2=FOO']}
     mock.assert_called_with(**expected_args)
+
+
+def test_pass_arbitrary_kwargs_from_yaml():
+    f = Flow.load_config(os.path.join(cur_dir, 'flow.yml'))
+    assert f._pod_nodes['pod1'].args.docker_kwargs == {'hello': 0, 'environment': ['VAR1=BAR', 'VAR2=FOO']}