Merge pull request #36 from Yoctol/use-grpclib

Use grpclib not aiorpc

Author: SoluMilken

.gitignore

@@ -96,5 +96,4 @@ dump.rdb
 .vscode/
 .pytest_cache/
 
-**/*_pb2*
 .fake-models

Pipfile

@@ -3,16 +3,13 @@ url = "https://pypi.org/simple"
 verify_ssl = true
 name = "pypi"
 
-
 [packages]
 serving-utils = {editable = true, path = "."}
 
-
 [dev-packages]
 pytest = "*"
 "flake8-config-yoctol" = ">=0.0.11"
 pytest-asyncio = "*"
 
-
 [requires]
 python_version = "3.6"

Pipfile.lock

@@ -102,3 +102,9 @@ make test
 ```
 make install-dev
 ```
+
+### Protos
+
+```
+python -m grpc_tools.protoc -I. --python_out=. --python_grpc_out=. --grpc_python_out=. serving_utils/protos/*.proto
+```

README.md

@@ -5,11 +5,12 @@
 from collections import namedtuple
 from concurrent.futures import ThreadPoolExecutor
 
-from aiogrpc import Channel
 import grpc
+from grpclib.client import Channel
 import tensorflow as tf
 
 from .protos import predict_pb2, prediction_service_pb2_grpc, list_models_pb2, list_models_pb2_grpc
+from .protos import prediction_service_grpc
 
 
 def copy_message(src, dst):
@@ -66,15 +67,18 @@ def __init__(
                 credentials=creds,
                 options=channel_options,
             )
-        self._async_channel = Channel(
-            self._channel,
-            loop=loop,
-            executor=executor,
-            standalone_pool_for_streaming=standalone_pool_for_streaming,
-        )
+
+        if loop is None:
+            loop = asyncio.get_event_loop()
+
+        split_addr = addr.split(':')
+        host = split_addr[0]
+        port = split_addr[1]
+        # TODO: better addr parsing, secure channel
+        self._async_channel = Channel(host, port, loop=loop)
 
         self._stub = prediction_service_pb2_grpc.PredictionServiceStub(self._channel)
-        self._async_stub = prediction_service_pb2_grpc.PredictionServiceStub(self._async_channel)
+        self._async_stub = prediction_service_grpc.PredictionServiceStub(self._async_channel)
         self._check_address_health()
 
     def _check_address_health(self):

serving_utils/client.py

@@ -0,0 +1,37 @@
+# Generated by the Protocol Buffers compiler. DO NOT EDIT!
+# source: serving_utils/protos/list_models.proto
+# plugin: grpclib.plugin.main
+import abc
+
+import grpclib.const
+import grpclib.client
+
+import serving_utils.protos.list_models_pb2
+
+
+class ListModelsBase(abc.ABC):
+
+    @abc.abstractmethod
+    async def ListModels(self, stream):
+        pass
+
+    def __mapping__(self):
+        return {
+            '/ListModels/ListModels': grpclib.const.Handler(
+                self.ListModels,
+                grpclib.const.Cardinality.UNARY_UNARY,
+                serving_utils.protos.list_models_pb2.ListModelsRequest,
+                serving_utils.protos.list_models_pb2.ListModelsResponse,
+            ),
+        }
+
+
+class ListModelsStub:
+
+    def __init__(self, channel: grpclib.client.Channel) -> None:
+        self.ListModels = grpclib.client.UnaryUnaryMethod(
+            channel,
+            '/ListModels/ListModels',
+            serving_utils.protos.list_models_pb2.ListModelsRequest,
+            serving_utils.protos.list_models_pb2.ListModelsResponse,
+        )

serving_utils/protos/list_models_grpc.py

@@ -0,0 +1,46 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+import grpc
+
+from serving_utils.protos import list_models_pb2 as serving__utils_dot_protos_dot_list__models__pb2
+
+
+class ListModelsStub(object):
+  # missing associated documentation comment in .proto file
+  pass
+
+  def __init__(self, channel):
+    """Constructor.
+
+    Args:
+      channel: A grpc.Channel.
+    """
+    self.ListModels = channel.unary_unary(
+        '/ListModels/ListModels',
+        request_serializer=serving__utils_dot_protos_dot_list__models__pb2.ListModelsRequest.SerializeToString,
+        response_deserializer=serving__utils_dot_protos_dot_list__models__pb2.ListModelsResponse.FromString,
+        )
+
+
+class ListModelsServicer(object):
+  # missing associated documentation comment in .proto file
+  pass
+
+  def ListModels(self, request, context):
+    # missing associated documentation comment in .proto file
+    pass
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details('Method not implemented!')
+    raise NotImplementedError('Method not implemented!')
+
+
+def add_ListModelsServicer_to_server(servicer, server):
+  rpc_method_handlers = {
+      'ListModels': grpc.unary_unary_rpc_method_handler(
+          servicer.ListModels,
+          request_deserializer=serving__utils_dot_protos_dot_list__models__pb2.ListModelsRequest.FromString,
+          response_serializer=serving__utils_dot_protos_dot_list__models__pb2.ListModelsResponse.SerializeToString,
+      ),
+  }
+  generic_handler = grpc.method_handlers_generic_handler(
+      'ListModels', rpc_method_handlers)
+  server.add_generic_rpc_handlers((generic_handler,))

serving_utils/protos/list_models_pb2.py

@@ -0,0 +1,3 @@
+# Generated by the Protocol Buffers compiler. DO NOT EDIT!
+# source: serving_utils/protos/model.proto
+# plugin: grpclib.plugin.main