add stream to stream rpc support

lovemefan · lovemefan · commit b3839e21b932 · 2023-09-19T17:26:02.000+08:00
diff --git a/README.md b/README.md
@@ -16,7 +16,7 @@
 > 深度学习模型在落地时需要提供高效快速交互接口，业务逻辑和深度模型解码通常运行在不同类型的机器上。
 Http 并不适合大量数据的交互，而RPC (Remote Procedure Call) 远程过程调用, 而RPC在TCP层实现。提高了开发效率，算法工程师可以不必花费更多精力放在具体的接口实现上，而是专注于算法优化上。
 
-tzrpc 框架基于google的 [grpc](https://github.com/grpc/) 实现，需要Python 3.7及以上
+tzrpc 框架基于google的 [grpc](https://github.com/grpc/) 实现，需要Python 3.7及以上， 支持流式传输！！！
 
 目前支持以下基础类型：
 
@@ -83,6 +83,11 @@ def send_bool(_bool: bool):
 def send_python_obj(data):
     return data
 
+@server.register(stream=True)
+def gumbel(num):
+    if num % 3 == 0:
+        yield f"number is {num}, you win"
+
 if __name__ == '__main__':
     server.run("localhost", 8000)
 ```
@@ -128,6 +133,11 @@ def send_bool(_bool: bool):
 def send_python_obj(data):
     return data
 
+@client.register(stream=True)
+def gumbel(num):
+    for i in range(num):
+        yield i
+
 if __name__ == '__main__':
     print(say_hello("lovemefan"))
     print(send_numpy_obj())
@@ -145,6 +155,10 @@ if __name__ == '__main__':
 
     python_obj = testOb("test_name", 20)
     print(send_python_obj(python_obj).__dict__)
+    
+    # 流式demo
+    for i in gumbel(10):
+        print(i)
 ```
 
 ### 客户端输出
@@ -168,4 +182,9 @@ False
 True
 
 {'name': 'test_name', 'age': 20}
+
+number is 0, you win
+number is 3, you win
+number is 6, you win
+number is 9, you win
 ```
diff --git a/setup.py b/setup.py
@@ -12,7 +12,7 @@
 
 setuptools.setup(
     name="tzrpc",
-    version="v0.0.3",
+    version="v0.0.4",
     author="Lovemefan, Yunnan Key Laboratory of Artificial Intelligence, "
     "Kunming University of Science and Technology, Kunming, Yunnan ",
     author_email="lovemefan@outlook.com",
diff --git a/tzrpc/client/client.py b/tzrpc/client/client.py
@@ -7,8 +7,6 @@
 import numbers
 import pickle
 from functools import partial
-from typing import Generator
-
 import grpc
 import numpy as np
 
@@ -50,11 +48,14 @@ def __init__(self, server_address: str, debug=False):
         if debug:
             logger.setLevel(logging.DEBUG)
         self.server_address = server_address
-        self.channel = grpc.insecure_channel(server_address,
-                                             options=[('grpc.max_send_message_length', MAX_MESSAGE_LENGTH),
-                                                      ('grpc.max_receive_message_length', MAX_MESSAGE_LENGTH),
-                                                      ('grpc.max_metadata_size', MAX_METADATA_SIZE)]
-                                             )
+        self.channel = grpc.insecure_channel(
+            server_address,
+            options=[
+                ("grpc.max_send_message_length", MAX_MESSAGE_LENGTH),
+                ("grpc.max_receive_message_length", MAX_MESSAGE_LENGTH),
+                ("grpc.max_metadata_size", MAX_METADATA_SIZE),
+            ],
+        )
 
     def register(self, func=None, stream=False):
         """
@@ -63,54 +64,69 @@ def register(self, func=None, stream=False):
             stream(bool): use stream mode
         :return:
         """
+
         # if return_type not in self.__type:
         #     raise ValueError(f"TZRPC return type only support {self.__type}")
         def decorate(_func):
             return partial(wrapper, _func)
 
         def wrapper(_func, *args, **kwargs):
             stub = toObjectStub(self.channel, _func.__name__)
+            # logger.debug(f"{self.channel._channel.name}-{_func.__name__}")
             result = _func(*args, **kwargs)
-            if isinstance(result, str):
-                request = String(text=result)
-                response = stub.toString(request).text
+            if not stream:
+                if isinstance(result, str):
+                    request = String(text=result)
+                    response = stub.toString(request).text
 
-            elif isinstance(result, np.ndarray):
-                request = numpy2protobuf(result)
-                response = protobuf2numpy(stub.toNdarray(request))
-            elif isinstance(result, bytes):
-                request = Bytes(data=result)
-                response = stub.toBytes(request).data
-            elif isinstance(result, bool):
-                request = Boolean(value=result)
-                response = bool(stub.toBoolean(request).value)
-            elif isinstance(result, numbers.Number):
-                if isinstance(result, int):
-                    request = Integer(value=result)
-                    response = stub.toInteger(request).value
-                elif isinstance(result, float):
-                    request = Double(value=result)
-                    response = stub.toDouble(request).value
+                elif isinstance(result, np.ndarray):
+                    request = numpy2protobuf(result)
+                    response = protobuf2numpy(stub.toNdarray(request))
+                elif isinstance(result, bytes):
+                    request = Bytes(data=result)
+                    response = stub.toBytes(request).data
+                elif isinstance(result, bool):
+                    request = Boolean(value=result)
+                    response = bool(stub.toBoolean(request).value)
+                elif isinstance(result, numbers.Number):
+                    if isinstance(result, int):
+                        request = Integer(value=result)
+                        response = stub.toInteger(request).value
+                    elif isinstance(result, float):
+                        request = Double(value=result)
+                        response = stub.toDouble(request).value
+                    else:
+                        logger.exception(f"Type of {type(result)} is not support")
+                elif self.tensor_type is not None and isinstance(
+                    result, self.tensor_type
+                ):
+                    request = numpy2protobuf(result.numpy())
+                    response = torch.from_numpy(
+                        protobuf2numpy(stub.toNdarray(request)).copy()
+                    )
                 else:
-                    logger.exception(f"Type of {type(result)} is not support")
-            elif self.tensor_type is not None and isinstance(result, self.tensor_type):
-                request = numpy2protobuf(result.numpy())
-                response = torch.from_numpy(
-                    protobuf2numpy(stub.toNdarray(request)).copy()
-                )
-            else:
-                logger.info(f"Data will serialized by pickle and send with bytes")
-                obj = pickle.dumps(result)
-                if stream:
-                    request = Bytes(data=b"STREAM"+obj)
-                    response = stub.toBytesStream(request)
-                else:
-                    request = Bytes(data=b"PICKLE"+obj)
+                    logger.info("Data will serialized by pickle and send with bytes")
+                    obj = pickle.dumps(result)
+                    request = Bytes(data=b"PICKLE" + obj)
                     response = pickle.loads(stub.toBytes(request).data)
+                    logger.debug(f"type of data loaded is {type(response)}")
+                return response
+            else:
+
+                def pack(data):
+                    for i in data:
+                        _obj = pickle.dumps(i)
+                        yield Bytes(data=b"PICKLE" + _obj)
 
-                logger.debug(f"type of data loaded is {type(response)}")
+                def unpack(data):
+                    for i in data:
+                        _data = i.data
+                        is_pickled = True if b"PICKLE" in _data[:6] else False
+                        if is_pickled:
+                            _data = pickle.loads(_data[6:])
+                        yield _data
 
-            return response
+                return unpack(stub.toBytesStream(pack(result)))
 
         if func is not None:
             return partial(wrapper, func)
diff --git a/tzrpc/decorator/rpc.py b/tzrpc/decorator/rpc.py
@@ -5,7 +5,6 @@
 # @File : rpc.py
 import pickle
 from functools import partial
-
 from tzrpc.proto.py.Boolean_pb2 import Boolean
 from tzrpc.proto.py.Bytes_pb2 import Bytes
 from tzrpc.proto.py.Number_pb2 import Double, Float, Integer
@@ -16,15 +15,14 @@
 
 servicers = []
 
-logger = get_logger(to_std=True, stdout_level="INFO", save_log_file=False)
+logger = get_logger(to_std=True, stdout_level="DEBUG", save_log_file=False)
 
 
 class RpcServicer:
     def __init__(self):
         pass
 
     def register(self, task=None, stream=False):
-
         if task is not None:
             _listener = Listener(task)
             # print(_listener)
@@ -119,20 +117,21 @@ def toNdarrays(self, request, context):
         response = ndarrays(ndarray=_ndarray)
         return response
 
-    def toBytesStream(self,  request, context):
-        logger.debug(f"Method toBytes({request}, {context}) called.")
-        data = request.data
-        is_pickled = True if b"PICKLE" in data[:6] else False
-        if is_pickled:
-            data = pickle.loads(data[6:])
-        result = self.task(data)
-        if is_pickled:
-            result = pickle.dumps(result)
+    def toBytesStream(self, request, context):
+        logger.debug(f"Method toBytesStream({request}, {context}) called.")
 
-        response = Bytes(data=result)
-        return response
+        for data in request:
+            data = data.data
+            is_pickled = True if b"PICKLE" in data[:6] else False
+            if is_pickled:
+                data = pickle.loads(data[6:])
+            result = self.task(data)
+
+            for i in result:
+                _result = pickle.dumps(i)
+                yield Bytes(data=b"PICKLE" + _result)
 
-    def toNdarrayStream(self,  request, context):
+    def toNdarrayStream(self, request, context):
         pass
 
     def toTensor(self, request, context):
