fix: properly load gateway arguments from YAML (#5664)

jina-ai · Feb 8, 2023 · 8117000 · 8117000
1 parent 96e990d
commit 8117000
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 0 deletions.
diff --git a/jina/jaml/parsers/gateway/legacy.py b/jina/jaml/parsers/gateway/legacy.py
@@ -26,6 +26,33 @@ def parse(
         cls._init_from_yaml = True
         # tmp_p = {kk: expand_env_var(vv) for kk, vv in data.get('with', {}).items()}
 
+        for key in {
+            'name',
+            'port',
+            'protocol',
+            'host',
+            'tracing',
+            'graph_description',
+            'graph_conditions',
+            'deployments_addresses',
+            'deployments_metadata',
+            'deployments_no_reduce',
+            'timeout_send',
+            'retries',
+            'compression',
+            'runtime_name',
+            'prefetch',
+            'meter',
+            'log_config',
+        }:
+            if runtime_args and not runtime_args.get(key) and data.get(key):
+                runtime_args[key] = data.get(key)
+        if runtime_args.get('default_port'):
+            yaml_port = data.get('port')
+            if isinstance(yaml_port, int):
+                yaml_port = [yaml_port]
+            runtime_args['port'] = yaml_port or runtime_args.get('port')
+
         obj = cls(
             **data.get('with', {}),
             metas=data.get('metas', {}),

diff --git a/jina/serve/runtimes/gateway/__init__.py b/jina/serve/runtimes/gateway/__init__.py
@@ -98,6 +98,7 @@ async def async_setup(self):
                 'aio_tracing_client_interceptors': self.aio_tracing_client_interceptors(),
                 'tracing_client_interceptor': self.tracing_client_interceptor(),
                 'log_config': self.args.log_config,
+                'default_port': getattr(self.args, 'default_port', False),
             },
             py_modules=self.args.py_modules,
             extra_search_paths=self.args.extra_search_paths,

diff --git a/tests/unit/test_yamlparser.py b/tests/unit/test_yamlparser.py
@@ -186,6 +186,10 @@ def test_load_gateway_external_success(yaml_file, gateway_name):
         assert gateway.arg1 == 'hello'
         assert gateway.arg2 == 'world'
         assert gateway.arg3 == 'default-arg3'
+        assert gateway.runtime_args.timeout_send == 10
+        assert gateway.runtime_args.retries == 10
+        assert gateway.runtime_args.compression == 'Deflate'
+        assert gateway.runtime_args.prefetch == 100
 
 
 @pytest.mark.parametrize(

diff --git a/tests/unit/yaml/test-custom-gateway.yml b/tests/unit/yaml/test-custom-gateway.yml
@@ -3,3 +3,7 @@ py_modules: dummy_gateway.py
 with:
   arg1: hello
   arg2: world
+timeout_send: 10
+retries: 10
+compression: Deflate
+prefetch: 100
diff --git a/tests/unit/yaml/test-fastapi-gateway.yml b/tests/unit/yaml/test-fastapi-gateway.yml
@@ -3,3 +3,7 @@ py_modules: dummy_fastapi_gateway.py
 with:
   arg1: hello
   arg2: world
+timeout_send: 10
+retries: 10
+compression: Deflate
+prefetch: 100