Repository: jina-ai/serve
Pull request: feat: add grpc streaming endpoint to worker and head runtimes
closes: https://github.com/jina-ai/jina/issues/5604

Problem statement:
Add the JinaRPC service to workerRuntime and HeadRuntime
Add the JinaRPC service to workerRuntime and HeadRuntime. This will allow jina client to communicate with both runtimes using `stream=True`.
Currently `stream=False` already works but not tested.
Add tests for client communicating with both runtimes in both cases

Hints from issue discussion:

Commit message: fix: await process data in head runtime

Patch to review:
diff --git a/jina/proto/jina.proto b/jina/proto/jina.proto
index f844eeffa4812..a021b39d6f1e1 100644
--- a/jina/proto/jina.proto
+++ b/jina/proto/jina.proto
@@ -160,7 +160,7 @@ service JinaSingleDataRequestRPC {
 }
 
 /**
- * jina Gateway gRPC service.
+ * jina streaming gRPC service.
  */
 service JinaRPC {
     // Pass in a Request and a filled Request with matches will be returned.
diff --git a/jina/serve/runtimes/head/__init__.py b/jina/serve/runtimes/head/__init__.py
index d7e1037cb7f86..b8ccd576f911a 100644
--- a/jina/serve/runtimes/head/__init__.py
+++ b/jina/serve/runtimes/head/__init__.py
@@ -4,7 +4,7 @@
 import os
 from abc import ABC
 from collections import defaultdict
-from typing import List
+from typing import TYPE_CHECKING, AsyncIterator, List
 
 import grpc
 from grpc_health.v1 import health, health_pb2, health_pb2_grpc
@@ -20,6 +20,9 @@
 from jina.serve.runtimes.helper import _get_grpc_server_options
 from jina.types.request.data import DataRequest, Response
 
+if TYPE_CHECKING:  # pragma: no cover
+    from jina.types.request import Request
+
 
 class HeadRuntime(AsyncNewLoopRuntime, ABC):
     """
@@ -134,6 +137,8 @@ async def async_setup(self):
             interceptors=self.aio_tracing_server_interceptors(),
         )
 
+        jina_pb2_grpc.add_JinaRPCServicer_to_server(self, self._grpc_server)
+
         jina_pb2_grpc.add_JinaSingleDataRequestRPCServicer_to_server(
             self, self._grpc_server
         )
@@ -143,6 +148,7 @@ async def async_setup(self):
         )
         jina_pb2_grpc.add_JinaInfoRPCServicer_to_server(self, self._grpc_server)
         service_names = (
+            jina_pb2.DESCRIPTOR.services_by_name['JinaRPC'].full_name,
             jina_pb2.DESCRIPTOR.services_by_name['JinaSingleDataRequestRPC'].full_name,
             jina_pb2.DESCRIPTOR.services_by_name['JinaDataRequestRPC'].full_name,
             jina_pb2.DESCRIPTOR.services_by_name['JinaDiscoverEndpointsRPC'].full_name,
@@ -318,3 +324,20 @@ async def _status(self, empty, context) -> jina_pb2.JinaInfoProto:
         for k, v in env_info.items():
             infoProto.envs[k] = str(v)
         return infoProto
+
+    async def stream(
+        self, request_iterator, context=None, *args, **kwargs
+    ) -> AsyncIterator['Request']:
+        """
+        stream requests from client iterator and stream responses back.
+
+        :param request_iterator: iterator of requests
+        :param context: context of the grpc call
+        :param args: positional arguments
+        :param kwargs: keyword arguments
+        :yield: responses to the request
+        """
+        async for request in request_iterator:
+            yield await self.process_data([request], context)
+
+    Call = stream
diff --git a/jina/serve/runtimes/worker/__init__.py b/jina/serve/runtimes/worker/__init__.py
index db5e8034e46ca..29d5b4dd338bf 100644
--- a/jina/serve/runtimes/worker/__init__.py
+++ b/jina/serve/runtimes/worker/__init__.py
@@ -2,7 +2,7 @@
 import asyncio
 import os
 from abc import ABC
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, AsyncIterator, List, Optional
 
 import grpc
 from grpc_health.v1 import health, health_pb2, health_pb2_grpc
@@ -21,6 +21,8 @@
 if TYPE_CHECKING:  # pragma: no cover
     from opentelemetry.propagate import Context
 
+    from jina.types.request import Request
+
 
 class WorkerRuntime(AsyncNewLoopRuntime, ABC):
     """Runtime procedure leveraging :class:`Grpclet` for sending DataRequests"""
@@ -121,6 +123,9 @@ async def _async_setup_grpc_server(self):
             interceptors=self.aio_tracing_server_interceptors(),
         )
 
+        print('adding JinaRPCServicer')
+        jina_pb2_grpc.add_JinaRPCServicer_to_server(self, self._grpc_server)
+
         jina_pb2_grpc.add_JinaSingleDataRequestRPCServicer_to_server(
             self, self._grpc_server
         )
@@ -131,6 +136,7 @@ async def _async_setup_grpc_server(self):
         )
         jina_pb2_grpc.add_JinaInfoRPCServicer_to_server(self, self._grpc_server)
         service_names = (
+            jina_pb2.DESCRIPTOR.services_by_name['JinaRPC'].full_name,
             jina_pb2.DESCRIPTOR.services_by_name['JinaSingleDataRequestRPC'].full_name,
             jina_pb2.DESCRIPTOR.services_by_name['JinaDataRequestRPC'].full_name,
             jina_pb2.DESCRIPTOR.services_by_name['JinaDiscoverEndpointsRPC'].full_name,
@@ -336,3 +342,20 @@ async def Watch(
         """
         self.logger.debug(f'Receive Watch request')
         return await self._health_servicer.Watch(request, context)
+
+    async def stream(
+        self, request_iterator, context=None, *args, **kwargs
+    ) -> AsyncIterator['Request']:
+        """
+        stream requests from client iterator and stream responses back.
+
+        :param request_iterator: iterator of requests
+        :param context: context of the grpc call
+        :param args: positional arguments
+        :param kwargs: keyword arguments
+        :yield: responses to the request
+        """
+        async for request in request_iterator:
+            yield await self.process_data([request], context)
+
+    Call = stream
diff --git a/tests/unit/serve/runtimes/head/test_head_runtime.py b/tests/unit/serve/runtimes/head/test_head_runtime.py
index a1183272ce953..ccc269ed71f2b 100644
--- a/tests/unit/serve/runtimes/head/test_head_runtime.py
+++ b/tests/unit/serve/runtimes/head/test_head_runtime.py
@@ -7,8 +7,8 @@
 
 import grpc
 import pytest
-
 from docarray import Document, DocumentArray
+
 from jina.clients.request import request_generator
 from jina.enums import PollingType
 from jina.proto import jina_pb2_grpc
@@ -20,7 +20,9 @@
 from tests.helper import _generate_pod_args
 
 
-def test_regular_data_case():
+@pytest.mark.parametrize('stream', [True, False])
+@pytest.mark.asyncio
+async def test_regular_data_case(stream):
     args = _generate_pod_args()
     args.polling = PollingType.ANY
     connection_list_dict = {0: [f'fake_ip:8080']}
@@ -31,12 +33,19 @@ def test_regular_data_case():
         f'{args.host}:{args.port}',
         options=GrpcConnectionPool.get_default_grpc_options(),
     ) as channel:
-        stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
-        response, call = stub.process_single_data.with_call(_create_test_data_message())
-
-    assert response
-    assert 'is-error' in dict(call.trailing_metadata())
-    assert len(response.docs) == 1
+        if stream:
+            stub = jina_pb2_grpc.JinaRPCStub(channel)
+            for resp in stub.Call(request_generator('/', DocumentArray.empty(3))):
+                assert resp
+                assert len(resp.docs) == 3
+        else:
+            stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
+            response, call = stub.process_single_data.with_call(
+                _create_test_data_message()
+            )
+            assert response
+            assert 'is-error' in dict(call.trailing_metadata())
+            assert len(response.docs) == 1
     assert not handle_queue.empty()
 
     _destroy_runtime(args, cancel_event, runtime_thread)
diff --git a/tests/unit/serve/runtimes/worker/test_worker_runtime.py b/tests/unit/serve/runtimes/worker/test_worker_runtime.py
index 3506b36b41131..318b870d78877 100644
--- a/tests/unit/serve/runtimes/worker/test_worker_runtime.py
+++ b/tests/unit/serve/runtimes/worker/test_worker_runtime.py
@@ -8,9 +8,8 @@
 import grpc
 import pytest
 import requests as req
-from jina import Document
 
-from jina import DocumentArray, Executor, requests
+from jina import Client, Document, DocumentArray, Executor, requests
 from jina.clients.request import request_generator
 from jina.proto import jina_pb2, jina_pb2_grpc
 from jina.serve.networking import GrpcConnectionPool
@@ -20,9 +19,16 @@
 from tests.helper import _generate_pod_args
 
 
+def _gen_test_data_message():
+    for _ in range(3):
+        yield _create_test_data_message()
+
+
 @pytest.mark.slow
 @pytest.mark.timeout(5)
-def test_worker_runtime():
+@pytest.mark.parametrize('stream', [True, False])
+@pytest.mark.asyncio
+async def test_worker_runtime(stream):
     args = _generate_pod_args()
 
     cancel_event = multiprocessing.Event()
@@ -49,14 +55,20 @@ def start_runtime(args, cancel_event):
         target,
         options=GrpcConnectionPool.get_default_grpc_options(),
     ) as channel:
-        stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
-        response, call = stub.process_single_data.with_call(_create_test_data_message())
+        if stream:
+            stub = jina_pb2_grpc.JinaRPCStub(channel)
+            for resp in stub.Call(request_generator('/', DocumentArray.empty(3))):
+                assert resp
+        else:
+            stub = jina_pb2_grpc.JinaSingleDataRequestRPCStub(channel)
+            response, call = stub.process_single_data.with_call(
+                _create_test_data_message()
+            )
+            assert response
 
     cancel_event.set()
     runtime_thread.join()
 
-    assert response
-
     assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')
 
 


Code review comments:
diff_hunk: @@ -121,6 +123,9 @@ async def _async_setup_grpc_server(self):
             interceptors=self.aio_tracing_server_interceptors(),
         )
 
+        print('adding JinaRPCServicer')
path: jina/serve/runtimes/worker/__init__.py
text: @user1:
remove print