feat: introduces sync gRPC server (#27)

ab93 · web-flow · commit 0fbecc13a2ae · 2022-11-22T17:13:26.000-05:00
* udf multi-threaded sync server
* udsink multi-threaded sync server
* udf and udsink async servers with migration threadpool
* convert Message to a dataclass type


Signed-off-by: Avik Basu &lt;avikbasu93@gmail.com&gt;
diff --git a/.codecov.yml b/.codecov.yml
@@ -3,11 +3,11 @@ coverage:
     project:
       default:
         target: auto
-        threshold: 1%
+        threshold: 3%
     patch:
       default:
         target: auto
-        threshold: 1%
+        threshold: 5%
 
 ignore:
 - "examples/"
diff --git a/.coveragerc b/.coveragerc
@@ -5,3 +5,9 @@ source = pynumaflow
 omit =
     pynumaflow/tests/*
     examples/*
+
+[report]
+exclude_lines =
+    def start
+    def start_async
+    def __serve_async
diff --git a/pynumaflow/function/_dtypes.py b/pynumaflow/function/_dtypes.py
@@ -1,4 +1,6 @@
+from dataclasses import dataclass
 from datetime import datetime
+from functools import partialmethod
 from typing import TypeVar, Type, List
 
 DROP = b"U+005C__DROP__"
@@ -9,41 +11,34 @@
 Ms = TypeVar("Ms", bound="Messages")
 
 
+@dataclass(frozen=True)
 class Message:
-    def __init__(self, key: str, value: bytes):
-        self._key = key or ""
-        self._value = value or b""
-
-    def __str__(self):
-        return str({self._key: self._value})
-
-    def __repr__(self):
-        return str(self)
+    """
+    Basic datatype for data passing to the next vertex/vertices.
 
-    @property
-    def key(self) -> str:
-        return self._key
+    Args:
+        key: string key for vertex;
+             special values are ALL (send to all), DROP (drop message)
+        value: data in bytes
+    """
 
-    @property
-    def value(self) -> bytes:
-        return self._value
+    key: str = ""
+    value: bytes = b""
 
     @classmethod
     def to_vtx(cls: Type[M], key: str, value: bytes) -> M:
+        """
+        Returns a Message object to send value to a vertex.
+        """
         return cls(key, value)
 
-    @classmethod
-    def to_all(cls: Type[M], value: bytes) -> M:
-        return cls(ALL, value)
-
-    @classmethod
-    def to_drop(cls: Type[M]) -> M:
-        return cls(DROP, b"")
+    to_all = partialmethod(to_vtx, ALL)
+    to_drop = partialmethod(to_vtx, DROP, b"")
 
 
 class Messages:
-    def __init__(self):
-        self._messages = []
+    def __init__(self, *messages: M):
+        self._messages = list(messages) or []
 
     def __str__(self):
         return str(self._messages)
diff --git a/pynumaflow/function/server.py b/pynumaflow/function/server.py
@@ -1,6 +1,8 @@
 import asyncio
 import logging
-from os import environ
+import multiprocessing
+import os
+from concurrent.futures import ThreadPoolExecutor
 from typing import Callable, Iterator
 
 import grpc
@@ -16,12 +18,14 @@
 from pynumaflow.function.generated import udfunction_pb2_grpc
 from pynumaflow.types import NumaflowServicerContext
 
-if environ.get("PYTHONDEBUG"):
+if os.getenv("PYTHONDEBUG"):
     logging.basicConfig(level=logging.DEBUG)
 
 _LOGGER = logging.getLogger(__name__)
 
 UDFMapCallable = Callable[[str, Datum], Messages]
+_PROCESS_COUNT = multiprocessing.cpu_count()
+MAX_THREADS = int(os.getenv("MAX_THREADS", 0)) or (_PROCESS_COUNT * 4)
 
 
 class UserDefinedFunctionServicer(udfunction_pb2_grpc.UserDefinedFunctionServicer):
@@ -33,15 +37,16 @@ class UserDefinedFunctionServicer(udfunction_pb2_grpc.UserDefinedFunctionService
         map_handler: Function callable following the type signature of UDFMapCallable
         sock_path: Path to the UNIX Domain Socket
         max_message_size: The max message size in bytes the server can receive and send
+        max_threads: The max number of threads to be spawned;
+                     defaults to number of processors x4
 
     Example invocation:
     >>> from pynumaflow.function import Messages, Message, Datum, UserDefinedFunctionServicer
     >>> def map_handler(key: str, datum: Datum) -> Messages:
     ...   val = datum.value
     ...   _ = datum.event_time
     ...   _ = datum.watermark
-    ...   messages = Messages()
-    ...   messages.append(Message.to_vtx(key, val))
+    ...   messages = Messages(Message.to_vtx(key, val))
     ...   return messages
     >>> grpc_server = UserDefinedFunctionServicer(map_handler)
     >>> grpc_server.start()
@@ -52,12 +57,19 @@ def __init__(
         map_handler: UDFMapCallable,
         sock_path=FUNCTION_SOCK_PATH,
         max_message_size=MAX_MESSAGE_SIZE,
+        max_threads=MAX_THREADS,
     ):
         self.__map_handler: UDFMapCallable = map_handler
         self.sock_path = f"unix://{sock_path}"
         self._max_message_size = max_message_size
+        self._max_threads = max_threads
         self._cleanup_coroutines = []
 
+        self._server_options = [
+            ("grpc.max_send_message_length", self._max_message_size),
+            ("grpc.max_receive_message_length", self._max_message_size),
+        ]
+
     def MapFn(
         self, request: udfunction_pb2.Datum, context: NumaflowServicerContext
     ) -> udfunction_pb2.DatumList:
@@ -112,35 +124,51 @@ def IsReady(
         """
         return udfunction_pb2.ReadyResponse(ready=True)
 
-    async def __serve(self) -> None:
-        server = grpc.aio.server(
-            options=[
-                ("grpc.max_send_message_length", self._max_message_size),
-                ("grpc.max_receive_message_length", self._max_message_size),
-            ]
-        )
+    async def __serve_async(self, server) -> None:
         udfunction_pb2_grpc.add_UserDefinedFunctionServicer_to_server(
             UserDefinedFunctionServicer(self.__map_handler), server
         )
         server.add_insecure_port(self.sock_path)
-        _LOGGER.info("Server listening on: %s", self.sock_path)
+        _LOGGER.info("GRPC Async Server listening on: %s", self.sock_path)
         await server.start()
 
         async def server_graceful_shutdown():
-            logging.info("Starting graceful shutdown...")
-            # Shuts down the server with 5 seconds of grace period. During the
-            # grace period, the server won't accept new connections and allow
-            # existing RPCs to continue within the grace period.
+            """
+            Shuts down the server with 5 seconds of grace period. During the
+            grace period, the server won't accept new connections and allow
+            existing RPCs to continue within the grace period.
+            """
+            _LOGGER.info("Starting graceful shutdown...")
             await server.stop(5)
 
         self._cleanup_coroutines.append(server_graceful_shutdown())
         await server.wait_for_termination()
 
-    def start(self) -> None:
-        """Starts the server on the given UNIX socket."""
+    def start_async(self) -> None:
+        """Starts the Async gRPC server on the given UNIX socket."""
+        server = grpc.aio.server(
+            ThreadPoolExecutor(max_workers=self._max_threads), options=self._server_options
+        )
         loop = asyncio.get_event_loop()
         try:
-            loop.run_until_complete(self.__serve())
+            loop.run_until_complete(self.__serve_async(server))
         finally:
             loop.run_until_complete(*self._cleanup_coroutines)
             loop.close()
+
+    def start(self) -> None:
+        """
+        Starts the gRPC server on the given UNIX socket with given max threads.
+        """
+        server = grpc.server(
+            ThreadPoolExecutor(max_workers=self._max_threads), options=self._server_options
+        )
+        udfunction_pb2_grpc.add_UserDefinedFunctionServicer_to_server(
+            UserDefinedFunctionServicer(self.__map_handler), server
+        )
+        server.add_insecure_port(self.sock_path)
+        server.start()
+        _LOGGER.info(
+            "GRPC Server listening on: %s with max threads: %s", self.sock_path, self._max_threads
+        )
+        server.wait_for_termination()
diff --git a/pynumaflow/sink/_dtypes.py b/pynumaflow/sink/_dtypes.py
@@ -6,7 +6,7 @@
 Rs = TypeVar("Rs", bound="Responses")
 
 
-@dataclass
+@dataclass(frozen=True)
 class Response:
     id: str
     success: bool
diff --git a/pynumaflow/sink/server.py b/pynumaflow/sink/server.py
@@ -1,6 +1,8 @@
 import asyncio
 import logging
-from os import environ
+import multiprocessing
+import os
+from concurrent.futures import ThreadPoolExecutor
 from typing import Callable, List
 
 import grpc
@@ -14,12 +16,14 @@
 from pynumaflow.sink.generated import udsink_pb2_grpc, udsink_pb2
 from pynumaflow.types import NumaflowServicerContext
 
-if environ.get("PYTHONDEBUG"):
+if os.getenv("PYTHONDEBUG"):
     logging.basicConfig(level=logging.DEBUG)
 
 _LOGGER = logging.getLogger(__name__)
 
 UDSinkCallable = Callable[[List[Datum]], Responses]
+_PROCESS_COUNT = multiprocessing.cpu_count()
+MAX_THREADS = int(os.getenv("MAX_THREADS", 0)) or (_PROCESS_COUNT * 4)
 
 
 class UserDefinedSinkServicer(udsink_pb2_grpc.UserDefinedSinkServicer):
@@ -31,14 +35,15 @@ class UserDefinedSinkServicer(udsink_pb2_grpc.UserDefinedSinkServicer):
         sink_handler: Function callable following the type signature of UDSinkCallable
         sock_path: Path to the UNIX Domain Socket
         max_message_size: The max message size in bytes the server can receive and send
+        max_threads: The max number of threads to be spawned;
+                     defaults to number of processors x 4
 
     Example invocation:
     >>> from typing import List
     >>> from pynumaflow.sink import Datum, Responses, Response, UserDefinedSinkServicer
     >>> def udsink_handler(datums: List[Datum]) -> Responses:
     ...   responses = Responses()
     ...   for msg in datums:
-    ...     print("User Defined Sink", msg)
     ...     responses.append(Response.as_success(msg.id))
     ...   return responses
     >>> grpc_server = UserDefinedSinkServicer(udsink_handler)
@@ -50,12 +55,19 @@ def __init__(
         sink_handler: UDSinkCallable,
         sock_path=SINK_SOCK_PATH,
         max_message_size=MAX_MESSAGE_SIZE,
+        max_threads=MAX_THREADS,
     ):
         self.__sink_handler: UDSinkCallable = sink_handler
         self.sock_path = f"unix://{sock_path}"
         self._max_message_size = max_message_size
+        self._max_threads = max_threads
         self._cleanup_coroutines = []
 
+        self._server_options = [
+            ("grpc.max_send_message_length", self._max_message_size),
+            ("grpc.max_receive_message_length", self._max_message_size),
+        ]
+
     def SinkFn(
         self, request: udsink_pb2.DatumList, context: NumaflowServicerContext
     ) -> udsink_pb2.ResponseList:
@@ -90,35 +102,51 @@ def IsReady(
         """
         return udsink_pb2.ReadyResponse(ready=True)
 
-    async def __serve(self) -> None:
+    async def __serve_async(self) -> None:
         server = grpc.aio.server(
-            options=[
-                ("grpc.max_send_message_length", self._max_message_size),
-                ("grpc.max_receive_message_length", self._max_message_size),
-            ]
+            ThreadPoolExecutor(max_workers=self._max_threads), options=self._server_options
         )
         udsink_pb2_grpc.add_UserDefinedSinkServicer_to_server(
             UserDefinedSinkServicer(self.__sink_handler), server
         )
         server.add_insecure_port(self.sock_path)
-        _LOGGER.info("Server listening on: %s", self.sock_path)
+        _LOGGER.info("GRPC Async Server listening on: %s", self.sock_path)
         await server.start()
 
         async def server_graceful_shutdown():
-            logging.info("Starting graceful shutdown...")
-            # Shuts down the server with 5 seconds of grace period. During the
-            # grace period, the server won't accept new connections and allow
-            # existing RPCs to continue within the grace period.
+            _LOGGER.info("Starting graceful shutdown...")
+            """
+            Shuts down the server with 5 seconds of grace period. During the
+            grace period, the server won't accept new connections and allow
+            existing RPCs to continue within the grace period.
             await server.stop(5)
+            """
 
         self._cleanup_coroutines.append(server_graceful_shutdown())
         await server.wait_for_termination()
 
-    def start(self) -> None:
-        """Starts the server on the given UNIX socket."""
+    def start_async(self) -> None:
+        """Starts the Async gRPC server on the given UNIX socket."""
         loop = asyncio.get_event_loop()
         try:
-            loop.run_until_complete(self.__serve())
+            loop.run_until_complete(self.__serve_async())
         finally:
             loop.run_until_complete(*self._cleanup_coroutines)
             loop.close()
+
+    def start(self) -> None:
+        """
+        Starts the gRPC server on the given UNIX socket with given max threads.
+        """
+        server = grpc.server(
+            ThreadPoolExecutor(max_workers=self._max_threads), options=self._server_options
+        )
+        udsink_pb2_grpc.add_UserDefinedSinkServicer_to_server(
+            UserDefinedSinkServicer(self.__sink_handler), server
+        )
+        server.add_insecure_port(self.sock_path)
+        server.start()
+        _LOGGER.info(
+            "GRPC Server listening on: %s with max threads: %s", self.sock_path, self._max_threads
+        )
+        server.wait_for_termination()
diff --git a/pynumaflow/tests/function/test_messages.py b/pynumaflow/tests/function/test_messages.py
