fix: fix custom gateway container pod (#5361)

jina-ai · Nov 8, 2022 · bef22c8 · bef22c8
1 parent c28a763
commit bef22c8
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 8 deletions.
diff --git a/jina/orchestrate/deployments/__init__.py b/jina/orchestrate/deployments/__init__.py
@@ -803,7 +803,9 @@ def _set_pod_args(args: Namespace) -> Dict[int, List[Namespace]]:
             for replica_id in range(replicas):
                 _args = copy.deepcopy(args)
                 _args.shard_id = shard_id
-                _args.pod_role = PodRoleType.WORKER
+                # for gateway pods, the pod role shouldn't be changed
+                if _args.pod_role != PodRoleType.GATEWAY:
+                    _args.pod_role = PodRoleType.WORKER
 
                 if cuda_device_map:
                     _args.env['CUDA_VISIBLE_DEVICES'] = str(cuda_device_map[replica_id])

diff --git a/jina/orchestrate/pods/__init__.py b/jina/orchestrate/pods/__init__.py
@@ -17,6 +17,8 @@
 
 __all__ = ['BasePod', 'Pod']
 
+from jina.serve.runtimes.gateway import GatewayRuntime
+
 
 def run(
     args: 'argparse.Namespace',
@@ -188,12 +190,21 @@ def _wait_for_ready_or_shutdown(self, timeout: Optional[float]):
         :param timeout: The time to wait before readiness or failure is determined
             .. # noqa: DAR201
         """
-        return AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
-            timeout=timeout,
-            ready_or_shutdown_event=self.ready_or_shutdown.event,
-            ctrl_address=self.runtime_ctrl_address,
-            timeout_ctrl=self._timeout_ctrl,
-        )
+        if self.args.pod_role == PodRoleType.GATEWAY:
+            return GatewayRuntime.wait_for_ready_or_shutdown(
+                timeout=timeout,
+                ready_or_shutdown_event=self.ready_or_shutdown.event,
+                ctrl_address=self.runtime_ctrl_address,
+                timeout_ctrl=self._timeout_ctrl,
+                protocol=self.args.protocol,
+            )
+        else:
+            return AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
+                timeout=timeout,
+                ready_or_shutdown_event=self.ready_or_shutdown.event,
+                ctrl_address=self.runtime_ctrl_address,
+                timeout_ctrl=self._timeout_ctrl,
+            )
 
     def _fail_start_timeout(self, timeout):
         """

diff --git a/tests/unit/orchestrate/pods/container/test_container_pod.py b/tests/unit/orchestrate/pods/container/test_container_pod.py
@@ -3,7 +3,7 @@
 
 import pytest
 
-from jina import __cache_path__
+from jina import Flow, __cache_path__
 from jina.excepts import RuntimeFailToStart
 from jina.helper import random_port
 from jina.orchestrate.pods.container import ContainerPod
@@ -281,3 +281,13 @@ def test_container_pod_custom_gateway(dummy_custom_gateway_docker_image_built):
     client = docker.from_env()
     containers = client.containers.list()
     assert container.id not in containers
+
+
+def test_container_pod_with_flow_custom_gateway(
+    dummy_custom_gateway_docker_image_built,
+):
+    flow = Flow(uses='docker://custom-gateway', protocol='http')
+    with flow:
+        _validate_dummy_custom_gateway_response(
+            flow.port, {'arg1': 'hello', 'arg2': 'world', 'arg3': 'default-arg3'}
+        )