feat(grpc): add StreamCall and decouple send and receive

Author: hanhxiao

gnes/client/cli.py

@@ -41,10 +41,10 @@ def __init__(self, args):
             stub = gnes_pb2_grpc.GnesRPCStub(channel)
 
             if args.mode == 'train':
-                resp = stub.RequestStreamCall(RequestGenerator.train(all_bytes, args.batch_size))
+                resp = list(stub.StreamCall(RequestGenerator.train(all_bytes, args.batch_size)))[-1]
                 print(resp)
             elif args.mode == 'index':
-                resp = stub.RequestStreamCall(RequestGenerator.train(all_bytes, args.batch_size))
+                resp = list(stub.StreamCall(RequestGenerator.train(all_bytes, args.batch_size)))[-1]
                 print(resp)
             elif args.mode == 'query':
                 for idx, q in enumerate(all_bytes):

gnes/client/http.py

@@ -91,7 +91,7 @@ async def init(loop):
             return srv
 
         def stub_call(req):
-            res_f = stub.RequestStreamCall(req)
+            res_f = list(stub.StreamCall(req))[-1]
             return json.loads(MessageToJson(res_f))
 
         with grpc.insecure_channel(

gnes/proto/gnes.proto

@@ -207,7 +207,7 @@ service GnesRPC {
     }
     rpc Call (Request) returns (Response) {
     }
-    rpc RequestStreamCall (stream Request) returns (Response) {
+    rpc StreamCall (stream Request) returns (stream Response) {
     }
 }
 

gnes/proto/gnes_pb2.py

@@ -34,8 +34,8 @@ def __init__(self, channel):
         request_serializer=gnes__pb2.Request.SerializeToString,
         response_deserializer=gnes__pb2.Response.FromString,
         )
-    self.RequestStreamCall = channel.stream_unary(
-        '/gnes.GnesRPC/RequestStreamCall',
+    self.StreamCall = channel.stream_stream(
+        '/gnes.GnesRPC/StreamCall',
         request_serializer=gnes__pb2.Request.SerializeToString,
         response_deserializer=gnes__pb2.Response.FromString,
         )
@@ -74,7 +74,7 @@ def Call(self, request, context):
     context.set_details('Method not implemented!')
     raise NotImplementedError('Method not implemented!')
 
-  def RequestStreamCall(self, request_iterator, context):
+  def StreamCall(self, request_iterator, context):
     # missing associated documentation comment in .proto file
     pass
     context.set_code(grpc.StatusCode.UNIMPLEMENTED)
@@ -104,8 +104,8 @@ def add_GnesRPCServicer_to_server(servicer, server):
           request_deserializer=gnes__pb2.Request.FromString,
           response_serializer=gnes__pb2.Response.SerializeToString,
       ),
-      'RequestStreamCall': grpc.stream_unary_rpc_method_handler(
-          servicer.RequestStreamCall,
+      'StreamCall': grpc.stream_stream_rpc_method_handler(
+          servicer.StreamCall,
           request_deserializer=gnes__pb2.Request.FromString,
           response_serializer=gnes__pb2.Response.SerializeToString,
       ),

gnes/proto/gnes_pb2_grpc.py

@@ -131,10 +131,14 @@ def Index(self, request, context):
     def Search(self, request, context):
         return self.Call(request, context)
 
-    def RequestStreamCall(self, request_iterator, context):
-        for request in request_iterator:
-            ret = self.Call(request, context)
-        return ret
+    def StreamCall(self, request_iterator, context):
+        num_result = 0
+        with self.zmq_context as zmq_client:
+            for request in request_iterator:
+                zmq_client.send_message(self.add_envelope(request, zmq_client), self.args.timeout)
+                num_result += 1
+            for _ in range(num_result):
+                yield self.remove_envelope(zmq_client.recv_message(self.args.timeout))
 
 
 class GRPCFrontend:

gnes/service/grpc.py

@@ -57,15 +57,10 @@ def test_grpc_frontend(self):
                          ('grpc.max_receive_message_length', 70 * 1024 * 1024)]) as channel:
             stub = gnes_pb2_grpc.GnesRPCStub(channel)
             with TimeContext('sync call'):  # about 5s
-                resp = stub.RequestStreamCall(RequestGenerator.train(self.all_bytes, 1))
+                resp = list(stub.StreamCall(RequestGenerator.train(self.all_bytes, 1)))[-1]
 
             self.assertEqual(resp.request_id, str(len(self.all_bytes)))  # idx start with 0, but +1 for final FLUSH
 
-            # test async calls
-            with TimeContext('async call'):  # immeidiately returns 0.001 s
-                resp = stub.RequestStreamCall.future(RequestGenerator.train(self.all_bytes, 1))
-            self.assertEqual(resp.result().request_id, str(len(self.all_bytes)))
-
     @unittest.mock.patch.dict(os.environ, {'http_proxy': '', 'https_proxy': ''})
     def test_async_block(self):
         args = set_grpc_frontend_parser().parse_args([
@@ -91,9 +86,19 @@ def test_async_block(self):
                 options=[('grpc.max_send_message_length', 70 * 1024 * 1024),
                          ('grpc.max_receive_message_length', 70 * 1024 * 1024)]) as channel:
             stub = gnes_pb2_grpc.GnesRPCStub(channel)
-            with TimeContext('sync call'):  # about 5s
-                resp = stub.RequestStreamCall.future(RequestGenerator.train(self.all_bytes, 1))
-
-            self.assertEqual(resp.result().request_id, str(len(self.all_bytes)))
-
-            self.assertEqual(resp.request_id, str(len(self.all_bytes2)))  # idx start with 0, but +1 for final FLUSH
+            id = 0
+            with TimeContext('non-blocking call'):  # about 26s = 32s (total) - 3*2s (overlap)
+                resp = stub.StreamCall(RequestGenerator.train(self.all_bytes2, 1))
+                for r in resp:
+                    self.assertEqual(r.request_id, str(id))
+                    id += 1
+
+            id = 0
+            with TimeContext('blocking call'):  # should be 32 s
+                for r in RequestGenerator.train(self.all_bytes2, 1):
+                    resp = stub.Call(r)
+                    self.assertEqual(resp.request_id, str(id))
+                    id += 1
+            # self.assertEqual(resp.result().request_id, str(len(self.all_bytes)))
+
+            # self.assertEqual(resp.request_id, str(len(self.all_bytes2)))  # idx start with 0, but +1 for final FLUSH