Repository: jina-ai/serve
Pull request: fix: start DataRequestHandler before grpc server
# Start DataRequestHandler before grpc server

## Description

We need to start the DataRequestHandler before the grpc Server. We may have Executors which are slow to init (downloading models etc) and we should not indicate to ready before the DataRequestHandler has been started.
This works atm with STATUS requests, but on K8s we are just doing plain TCP connects to the gRPC server to check if the pod is ready. These connects should fail while the Excutor is still initializing.

Closes #4570 ([issue](https://github.com/jina-ai/jina/issues/4570))

Problem statement:
WorkerRuntime starts gRPC Server before DataRequestHandler is ready
The WorkerRuntime calls `super().init()` before `DataRequestHandler(args, self.logger)`. The super init will in turn call `async_setup()`, which will initialize the grpc server. So this server will be ready BEFORE the DataRequestHandler.

This is bad, because a STATUS request will succeed signalling the runtime being ready when it may in fact not be. DataRequests will get stuck though.

Hints from issue discussion:

Commit message: fix: start DataRequestHandler before grpc server

Patch to review:
diff --git a/jina/serve/runtimes/worker/__init__.py b/jina/serve/runtimes/worker/__init__.py
index 5e73509f2dbc7..996e1c4e8fad6 100644
--- a/jina/serve/runtimes/worker/__init__.py
+++ b/jina/serve/runtimes/worker/__init__.py
@@ -32,13 +32,16 @@ def __init__(
         """
         super().__init__(args, cancel_event, **kwargs)
 
-        # Keep this initialization order, otherwise readiness check is not valid
-        self._data_request_handler = DataRequestHandler(args, self.logger)
-
     async def async_setup(self):
         """
-        Wait for the GRPC server to start
+        Start the DataRequestHandler and wait for the GRPC server to start
         """
+
+        # Keep this initialization order
+        # otherwise readiness check is not valid
+        # The DataRequestHandler needs to be started BEFORE the grpc server
+        self._data_request_handler = DataRequestHandler(self.args, self.logger)
+
         self._grpc_server = grpc.aio.server(
             options=[
                 ('grpc.max_send_message_length', -1),
diff --git a/tests/unit/serve/runtimes/worker/test_worker_runtime.py b/tests/unit/serve/runtimes/worker/test_worker_runtime.py
index 1125f49c0c00e..ccce1878dbd59 100644
--- a/tests/unit/serve/runtimes/worker/test_worker_runtime.py
+++ b/tests/unit/serve/runtimes/worker/test_worker_runtime.py
@@ -1,24 +1,23 @@
 import asyncio
 import multiprocessing
 import os
+import socket
 import time
 from multiprocessing import Process
 from threading import Event
 
 import grpc
 import pytest
-
 from docarray import Document
+
 from jina import DocumentArray, Executor, requests
 from jina.clients.request import request_generator
 from jina.parsers import set_pod_parser
+from jina.proto import jina_pb2, jina_pb2_grpc
 from jina.serve.networking import GrpcConnectionPool
 from jina.serve.runtimes.asyncio import AsyncNewLoopRuntime
-from jina.serve.runtimes.request_handlers.data_request_handler import (
-    DataRequestHandler,
-)
+from jina.serve.runtimes.request_handlers.data_request_handler import DataRequestHandler
 from jina.serve.runtimes.worker import WorkerRuntime
-from jina.proto import jina_pb2_grpc, jina_pb2
 
 
 @pytest.mark.slow
@@ -275,5 +274,69 @@ async def task_wrapper(adress, messages_received):
     assert not WorkerRuntime.is_ready(f'{args.host}:{args.port}')
 
 
+class SlowInitExecutor(Executor):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        time.sleep(5.0)
+
+    @requests
+    def foo(self, docs, **kwargs):
+        return docs
+
+
+@pytest.mark.timeout(10)
+@pytest.mark.asyncio
+async def test_worker_runtime_slow_init_exec():
+    args = set_pod_parser().parse_args(['--uses', 'SlowInitExecutor'])
+
+    cancel_event = multiprocessing.Event()
+
+    def start_runtime(args, cancel_event):
+        with WorkerRuntime(args, cancel_event) as runtime:
+            runtime.run_forever()
+
+    runtime_thread = Process(
+        target=start_runtime,
+        args=(args, cancel_event),
+        daemon=True,
+    )
+    runtime_started = time.time()
+    runtime_thread.start()
+
+    # wait a bit to the worker runtime has a chance to finish some things, but not the Executor init (5 secs)
+    time.sleep(1.0)
+
+    # try to connect a TCP socket to the gRPC server
+    # this should only succeed after the Executor is ready, which should be after 5 seconds
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    connected = False
+    while not connected:
+        try:
+            s.connect((args.host, args.port))
+            connected = True
+        except ConnectionRefusedError:
+            time.sleep(0.2)
+
+    # Executor sleeps 5 seconds, so at least 5 seconds need to have elapsed here
+    assert time.time() - runtime_started > 5.0
+
+    assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
+        timeout=3.0,
+        ctrl_address=f'{args.host}:{args.port}',
+        ready_or_shutdown_event=Event(),
+    )
+
+    result = await GrpcConnectionPool.send_request_async(
+        _create_test_data_message(), f'{args.host}:{args.port}', timeout=1.0
+    )
+
+    assert len(result.docs) == 1
+
+    cancel_event.set()
+    runtime_thread.join()
+
+    assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')
+
+
 def _create_test_data_message(counter=0):
     return list(request_generator('/', DocumentArray([Document(text=str(counter))])))[0]


Code review comments:
diff_hunk: @@ -275,5 +274,69 @@ async def task_wrapper(adress, messages_received):
     assert not WorkerRuntime.is_ready(f'{args.host}:{args.port}')
 
 
+class SlowInitExecutor(Executor):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        time.sleep(5.0)
+
+    @requests
+    def foo(self, docs, **kwargs):
+        return docs
+
+
+@pytest.mark.timeout(10)
+@pytest.mark.asyncio
+async def test_worker_runtime_slow_init_exec():
+    args = set_pod_parser().parse_args(['--uses', 'SlowInitExecutor'])
+
+    cancel_event = multiprocessing.Event()
+
+    def start_runtime(args, cancel_event):
+        with WorkerRuntime(args, cancel_event) as runtime:
+            runtime.run_forever()
+
+    runtime_thread = Process(
+        target=start_runtime,
+        args=(args, cancel_event),
+        daemon=True,
+    )
+    runtime_started = time.time()
+    runtime_thread.start()
+
+    # wait a bit to the worker runtime has a chance to finish some things, but not the Executor init (5 secs)
+    time.sleep(1.0)
+
+    # try to connect a TCP socket to the gRPC server
+    # this should only succeed after the Executor is ready, which should be after 5 seconds
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    connected = False
+    while not connected:
+        try:
+            s.connect((args.host, args.port))
+            connected = True
+        except ConnectionRefusedError:
+            time.sleep(0.2)
path: tests/unit/serve/runtimes/worker/test_worker_runtime.py
text: @user1:
I had some issue with pytest when I don't close the socket. Here it might not change anything but it might prevent future but in the test maybe
```suggestion
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        connected = False
        while not connected:
            try:
                s.connect((args.host, args.port))
                connected = True
            except ConnectionRefusedError:
                time.sleep(0.2)
```

@author:
good point, added it!