feat: add startupProbe and replace readiness probe with liveness probe

.github/workflows/cd.yml

@@ -271,7 +271,7 @@ jobs:
           export LINKERD2_VERSION=stable-2.11.4
           curl --proto '=https' --tlsv1.2 -sSfL https://run.linkerd.io/install | sh
           pytest -v -s --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml ./tests/k8s/test_k8s.py ./tests/k8s/test_graceful_request_handling.py
-        timeout-minutes: 30
+        timeout-minutes: 45
         env:
           JINA_K8S_USE_TEST_PIP: 1
       - name: Check codecov file

.github/workflows/ci.yml

@@ -194,7 +194,7 @@ jobs:
           export LINKERD2_VERSION=stable-2.11.4
           curl --proto '=https' --tlsv1.2 -sSfL https://run.linkerd.io/install | sh
           pytest -v -s --suppress-no-test-exit-code --force-flaky --min-passes 1 --max-runs 5 --cov=jina --cov-report=xml ./tests/k8s/test_k8s.py ./tests/k8s/test_graceful_request_handling.py
-        timeout-minutes: 30
+        timeout-minutes: 45
         env:
           JINA_K8S_USE_TEST_PIP: 1
       - name: Check codecov file

jina/checker.py

@@ -31,12 +31,13 @@ def __init__(self, args: 'argparse.Namespace'):
                 ) as tc:
                     if args.target == 'executor':
                         hostname, port, protocol, _ = parse_host_scheme(args.host)
-                        r = WorkerRuntime.is_ready(f'{hostname}:{port}')
+                        r = WorkerRuntime.is_ready(ctrl_address=f'{hostname}:{port}', timeout=args.timeout)
                     elif args.target == 'gateway':
                         hostname, port, protocol, _ = parse_host_scheme(args.host)
                         r = GatewayRuntime.is_ready(
                             f'{hostname}:{port}',
                             protocol=GatewayProtocolType.from_string(protocol),
+                            timeout=args.timeout
                         )
                     elif args.target == 'flow':
                         r = Client(host=args.host).is_flow_ready(timeout=args.timeout)

jina/serve/runtimes/worker/__init__.py

@@ -1,6 +1,7 @@
 import argparse
 from abc import ABC
 from typing import TYPE_CHECKING, List, Optional
+from concurrent.futures import ThreadPoolExecutor
 
 import grpc
 from grpc_health.v1 import health, health_pb2, health_pb2_grpc
@@ -32,7 +33,7 @@ def __init__(
         :param args: args from CLI
         :param kwargs: keyword args
         """
-        self._health_servicer = health.aio.HealthServicer()
+        self._health_servicer = health.HealthServicer(experimental_thread_pool=ThreadPoolExecutor(1))
         super().__init__(args, **kwargs)
 
     async def async_setup(self):
@@ -140,7 +141,7 @@ async def _async_setup_grpc_server(self):
         )
 
         for service in service_names:
-            await self._health_servicer.set(
+            self._health_servicer.set(
                 service, health_pb2.HealthCheckResponse.SERVING
             )
         reflection.enable_server_reflection(service_names, self._grpc_server)
@@ -164,7 +165,7 @@ async def async_cancel(self):
 
     async def async_teardown(self):
         """Close the data request handler"""
-        await self._health_servicer.enter_graceful_shutdown()
+        self._health_servicer.enter_graceful_shutdown()
         await self.async_cancel()
         self._request_handler.close()
 

tests/k8s/slow-process-executor/debug_executor.py

@@ -5,15 +5,13 @@
 
 
 class SlowProcessExecutor(Executor):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, time_sleep=1.0, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        from jina.logging.logger import JinaLogger
-
-        self.logger = JinaLogger(self.__class__.__name__)
+        self.time_sleep = time_sleep
 
     @requests
     def process(self, docs: DocumentArray, *args, **kwargs):
-        time.sleep(1.0)
+        time.sleep(self.time_sleep)
         for doc in docs:
             doc.tags['replica_uid'] = os.environ['POD_UID']
             doc.tags['time'] = time.time()

tests/k8s/test_k8s.py

@@ -1008,3 +1008,49 @@ async def test_flow_with_stateful_executor(
 
     assert len(resp) == 1
     assert resp[0].parameters == {'__results__': {'statefulexecutor': {'length': 10.0}}}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'docker_images', [['slow-process-executor', 'jinaai/jina']], indirect=True
+)
+async def test_slow_executor_readinessProbe_works(docker_images, tmpdir, logger):
+    dump_path = os.path.join(str(tmpdir), 'test-flow-slow-process-executor')
+    namespace = f'test-flow-slow-process-executor'.lower()
+    flow = Flow(name='test-flow-slow-process-executor',).add(
+        name='slow_process_executor',
+        uses=f'docker://{docker_images[0]}',
+        uses_with={'time_sleep': 200},
+        replicas=2,
+    )
+
+    flow.to_kubernetes_yaml(dump_path, k8s_namespace=namespace)
+
+    from kubernetes import client
+
+    api_client = client.ApiClient()
+    core_client = client.CoreV1Api(api_client=api_client)
+    app_client = client.AppsV1Api(api_client=api_client)
+    await create_all_flow_deployments_and_wait_ready(
+        dump_path,
+        namespace=namespace,
+        api_client=api_client,
+        app_client=app_client,
+        core_client=core_client,
+        deployment_replicas_expected={
+            'gateway': 1,
+            'slow-process-executor': 2,
+        },
+        logger=logger
+    )
+
+    resp = await run_test(
+        flow=flow,
+        namespace=namespace,
+        core_client=core_client,
+        n_docs=10,
+        request_size=1,
+        endpoint='/',
+    )
+
+    assert len(resp) == 10