fix: run health servicer in thread to allow readinessProbe while req

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

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()