Repository: jina-ai/serve
Pull request: feat: add startupProbe and replace readiness probe with liveness probe
<!---Decomposing the complex issue into subtasks can help you build it step-by-step. Thanks for your pull request! :rocket: --->
<!---We know that dev life is hectic, but **please provide a (brief) description** of what your PR does, and how it does it. **Otherwise, your PR cannot be reviewed!** --->
<!---This policy was agreed upon in a past company retro, and makes everyone's life a little easier. Thanks for your collaboration!--->

**Goals:**
Now, while Requests are being processed, the `event loop` of Worker is blocked and `readinessProbe` can start failing if Requests block the event loop too long (Executor methods taking long)

Resolves #5470 

Problem statement:
feat: add start up probe and replace readiness probe with liveness probe
**Describe the feature**
<!-- A clear and concise description of what the feature is. -->
Start up probe is better at handling the initial start up conditions of the application rather than increasing the `initialDelaySeconds` of the liveness or readiness probe.

**Your proposal**
<!-- copy past your code/pull request link -->
The start up probe will use the `jina ping` command to report the first container ready status.

---
<!-- Optional, but really help us locate the problem faster -->

**Environment**
<!-- Run `jina --version-full` and copy paste the output here -->

**Screenshots**
<!-- If applicable, add screenshots to help explain your problem. -->

Hints from issue discussion:

Commit message: test: add test with slow process

Signed-off-by: Joan Fontanals Martinez <joan.martinez@jina.ai>

Patch to review:
diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml
index f8171b42a200e..01a6f02928e70 100644
--- a/.github/workflows/cd.yml
+++ b/.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
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 277389b223313..0633453f39d9f 100644
--- a/.github/workflows/ci.yml
+++ b/.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
diff --git a/jina/checker.py b/jina/checker.py
index 25605d5a0bc00..6805d1480a8c4 100644
--- a/jina/checker.py
+++ b/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)
diff --git a/jina/serve/runtimes/worker/__init__.py b/jina/serve/runtimes/worker/__init__.py
index 4d58292e6d60f..d654627794153 100644
--- a/jina/serve/runtimes/worker/__init__.py
+++ b/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()
 
diff --git a/tests/k8s/slow-process-executor/debug_executor.py b/tests/k8s/slow-process-executor/debug_executor.py
index fb4c4e7068fda..c52d962449ffd 100644
--- a/tests/k8s/slow-process-executor/debug_executor.py
+++ b/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()
diff --git a/tests/k8s/test_k8s.py b/tests/k8s/test_k8s.py
index fd9db6f541c36..2f3b26b0a5767 100644
--- a/tests/k8s/test_k8s.py
+++ b/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


Code review comments:
diff_hunk: @@ -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))
path: jina/serve/runtimes/worker/__init__.py
text: @user1:
Removing the `aio.HealthServicer` will break the docker + OpenTelemetry fix that I made earlier this week. If we must use the regular `HealthServicer`, we need to adapt the tracing server interceptor to support both any type of rpc.

@author:
what was the reason of this? can u make a proposal of fix to this branch?

@user1:
I will create a separate pr to fix the tracing server interceptor.

@user1:
The [PR](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1245) for grpc aio support has been merged and released so we can remove our internal dependency in another pr.