feat: introduce handshake to client and gRPC server (#89)

Signed-off-by: Sidhant Kohli <[email protected]>

Author: kohlisid

examples/function/multiproc_map/README.md

@@ -15,13 +15,11 @@ processes to bind to the same port.
 
 To enable multiprocessing mode 
 
-1) Set the env flag `MAP_MULTIPROC="true"` for the numa container
-
-2) Start the multiproc server in the UDF using the following command 
+1) Start the multiproc server in the UDF using the following command 
 ```python
 if __name__ == "__main__":
     grpc_server = MultiProcServer(map_handler=my_handler)
     grpc_server.start()
 ```
-3) Set the ENV var value `NUM_CPU_MULTIPROC="n"` for the UDF and numa container,
-to set the value of the number of processes to be created. 
+2) Set the ENV var value `NUM_CPU_MULTIPROC="n"` for the UDF container,
+to set the value of the number of server instances (one for each subprocess) to be created.

examples/function/multiproc_map/pipeline.yaml

@@ -32,10 +32,6 @@ spec:
         env:
           - name: NUMAFLOW_DEBUG
             value: "true" # DO NOT forget the double quotes!!!
-          - name: MAP_MULTIPROC
-            value: "true" # DO NOT forget the double quotes!!!
-          - name: NUM_CPU_MULTIPROC
-            value: "2"  # DO NOT forget the double quotes!!!
 
     - name: out
       sink:

pynumaflow/function/async_server.py

@@ -23,6 +23,8 @@
 from pynumaflow.function.proto import udfunction_pb2
 from pynumaflow.function.proto import udfunction_pb2_grpc
 from pynumaflow.types import NumaflowServicerContext
+from pynumaflow.info.server import get_sdk_version, write as info_server_write
+from pynumaflow.info.types import ServerInfo, Protocol, Language, SERVER_INFO_FILE_PATH
 
 _LOGGER = setup_logging(__name__)
 if os.getenv("PYTHONDEBUG"):
@@ -306,6 +308,12 @@ async def __serve_async(self, server) -> None:
         server.add_insecure_port(self.sock_path)
         _LOGGER.info("GRPC Async Server listening on: %s", self.sock_path)
         await server.start()
+        serv_info = ServerInfo(
+            protocol=Protocol.UDS,
+            language=Language.PYTHON,
+            version=get_sdk_version(),
+        )
+        info_server_write(server_info=serv_info, info_file=SERVER_INFO_FILE_PATH)
 
         async def server_graceful_shutdown():
             """

pynumaflow/function/multiproc_server.py

@@ -21,6 +21,18 @@
 from pynumaflow.function.proto import udfunction_pb2
 from pynumaflow.function.proto import udfunction_pb2_grpc
 from pynumaflow.types import NumaflowServicerContext
+from pynumaflow.info.server import (
+    get_sdk_version,
+    write as info_server_write,
+    get_metadata_env,
+)
+from pynumaflow.info.types import (
+    ServerInfo,
+    Protocol,
+    Language,
+    SERVER_INFO_FILE_PATH,
+    METADATA_ENVS,
+)
 
 _LOGGER = setup_logging(__name__)
 if os.getenv("PYTHONDEBUG"):
@@ -29,8 +41,6 @@
 UDFMapCallable = Callable[[List[str], Datum], Messages]
 UDFMapTCallable = Callable[[List[str], Datum], MessageTs]
 UDFReduceCallable = Callable[[List[str], AsyncIterable[Datum], Metadata], Messages]
-_PROCESS_COUNT = int(os.getenv("NUM_CPU_MULTIPROC", multiprocessing.cpu_count()))
-MAX_THREADS = int(os.getenv("MAX_THREADS", 0)) or (_PROCESS_COUNT * 4)
 
 
 class MultiProcServer(udfunction_pb2_grpc.UserDefinedFunctionServicer):
@@ -84,7 +94,6 @@ def __init__(
         reduce_handler: UDFReduceCallable = None,
         sock_path=MULTIPROC_FUNCTION_SOCK_PORT,
         max_message_size=MAX_MESSAGE_SIZE,
-        max_threads=MAX_THREADS,
     ):
         if not (map_handler or mapt_handler or reduce_handler):
             raise ValueError("Require a valid map/mapt handler and/or a valid reduce handler.")
@@ -93,7 +102,6 @@ def __init__(
         self.__mapt_handler: UDFMapTCallable = mapt_handler
         self.__reduce_handler: UDFReduceCallable = reduce_handler
         self._max_message_size = max_message_size
-        self._max_threads = max_threads
         self.cleanup_coroutines = []
         # Collection for storing strong references to all running tasks.
         # Event loop only keeps a weak reference, which can cause it to
@@ -107,8 +115,10 @@ def __init__(
             ("grpc.so_reuseaddr", 1),
         ]
         self._sock_path = sock_path
-        self._process_count = int(os.getenv("NUM_CPU_MULTIPROC", multiprocessing.cpu_count()))
-        self._thread_concurrency = MAX_THREADS
+        self._process_count = int(
+            os.getenv("NUM_CPU_MULTIPROC") or os.getenv("NUMAFLOW_CPU_LIMIT", 1)
+        )
+        self._thread_concurrency = int(os.getenv("MAX_THREADS", 0)) or (self._process_count * 4)
 
     def MapFn(
         self, request: udfunction_pb2.DatumRequest, context: NumaflowServicerContext
@@ -229,6 +239,16 @@ def _run_server(self, bind_address):
         udfunction_pb2_grpc.add_UserDefinedFunctionServicer_to_server(self, server)
         server.add_insecure_port(bind_address)
         server.start()
+        serv_info = ServerInfo(
+            protocol=Protocol.TCP,
+            language=Language.PYTHON,
+            version=get_sdk_version(),
+            metadata=get_metadata_env(envs=METADATA_ENVS),
+        )
+        # Overwrite the CPU_LIMIT metadata using user input
+        serv_info.metadata["CPU_LIMIT"] = str(self._process_count)
+        info_server_write(server_info=serv_info, info_file=SERVER_INFO_FILE_PATH)
+
         _LOGGER.info("GRPC Multi-Processor Server listening on: %s %d", bind_address, os.getpid())
         server.wait_for_termination()
 

pynumaflow/function/server.py

@@ -7,6 +7,8 @@
 import grpc
 from google.protobuf import empty_pb2 as _empty_pb2
 from google.protobuf import timestamp_pb2 as _timestamp_pb2
+from pynumaflow.info.server import get_sdk_version, write as info_server_write
+from pynumaflow.info.types import ServerInfo, Protocol, Language, SERVER_INFO_FILE_PATH
 
 from pynumaflow import setup_logging
 from pynumaflow._constants import (
@@ -219,6 +221,12 @@ def start(self) -> None:
         udfunction_pb2_grpc.add_UserDefinedFunctionServicer_to_server(self, server)
         server.add_insecure_port(self.sock_path)
         server.start()
+        serv_info = ServerInfo(
+            protocol=Protocol.UDS,
+            language=Language.PYTHON,
+            version=get_sdk_version(),
+        )
+        info_server_write(server_info=serv_info, info_file=SERVER_INFO_FILE_PATH)
         _LOGGER.info(
             "GRPC Server listening on: %s with max threads: %s", self.sock_path, self._max_threads
         )

pynumaflow/info/__init__.py

@@ -0,0 +1,59 @@
+import os
+from importlib.metadata import version
+from typing import Any
+
+from pynumaflow import setup_logging
+from pynumaflow.info.types import ServerInfo, EOF
+import json
+import logging
+
+_LOGGER = setup_logging(__name__)
+if os.getenv("PYTHONDEBUG"):
+    _LOGGER.setLevel(logging.DEBUG)
+
+
+def get_sdk_version() -> str:
+    """
+    Return the pynumaflow SDK version
+    """
+    try:
+        return version("pynumaflow")
+    except Exception as e:
+        # Adding this to handle the case for local test/CI where pynumaflow
+        # will not be installed as a package
+        _LOGGER.error("Could not read SDK version %r", e, exc_info=True)
+        return ""
+
+
+def write(server_info: ServerInfo, info_file: str):
+    """
+    Write the ServerInfo to a file , shared with the client (numa container).
+
+    args:
+        serv: The ServerInfo object to be shared
+        info_file: the shared file path
+    """
+    try:
+        data = server_info.__dict__
+        with open(info_file, "w+") as f:
+            json.dump(data, f, ensure_ascii=False)
+            f.write(EOF)
+    except Exception as err:
+        _LOGGER.critical("Could not write data to Info-Server %r", err, exc_info=True)
+        raise err
+
+
+def get_metadata_env(envs: list[tuple[str, str]]) -> dict[str, Any]:
+    """
+    Extract the environment var value from the provided list,
+    and assign them to the given key in the metadata
+
+    args:
+        envs: List of tuples (key, env_var)
+    """
+    meta = {}
+    for key, val in envs:
+        res = os.getenv(val, None)
+        if res:
+            meta[key] = res
+    return meta

pynumaflow/info/server.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass, field
+from enum import Enum
+
+# Constants for using in the info-server
+# Need to keep consistent with all SDKs and client
+SERVER_INFO_FILE_PATH = "/var/run/numaflow/server-info"
+EOF = "U+005C__END__"
+
+# Env variables to be passed in the info server metadata.
+# These need to be accessed in the client using the same key.
+# Format - (key, env_var)
+METADATA_ENVS = [("CPU_LIMIT", "NUMAFLOW_CPU_LIMIT")]
+
+
+class Protocol(str, Enum):
+    """
+    Enumerate grpc server connection protocol.
+    """
+
+    UDS = "uds"
+    TCP = "tcp"
+
+
+class Language(str, Enum):
+    """
+    Enumerate Numaflow SDK language.
+    """
+
+    GO = "go"
+    PYTHON = "python"
+    JAVA = "java"
+
+
+@dataclass
+class ServerInfo:
+    """
+    ServerInfo is used for the gRPC server to provide the information such as protocol,
+    sdk version, language, metadata to the client.
+    Args:
+        protocol: Protocol to use (UDS or TCP)
+        language: Language used by the server(Python, Golang, Java)
+        version: Numaflow sdk version used by the server
+        metadata: Any additional information to be provided (env vars)
+    """
+
+    protocol: Protocol
+    language: Language
+    version: str
+    metadata: dict = field(default_factory=dict)

pynumaflow/info/types.py

@@ -11,6 +11,8 @@
     SINK_SOCK_PATH,
     MAX_MESSAGE_SIZE,
 )
+from pynumaflow.info.server import get_sdk_version, write as info_server_write
+from pynumaflow.info.types import ServerInfo, Protocol, Language, SERVER_INFO_FILE_PATH
 from pynumaflow.sink import Responses, Datum, Response
 from pynumaflow.sink.proto import udsink_pb2_grpc, udsink_pb2
 from pynumaflow.types import NumaflowServicerContext
@@ -128,6 +130,12 @@ async def __serve_async(self, server) -> None:
         server.add_insecure_port(self.sock_path)
         _LOGGER.info("GRPC Async Server listening on: %s", self.sock_path)
         await server.start()
+        serv_info = ServerInfo(
+            protocol=Protocol.UDS,
+            language=Language.PYTHON,
+            version=get_sdk_version(),
+        )
+        info_server_write(server_info=serv_info, info_file=SERVER_INFO_FILE_PATH)
 
         async def server_graceful_shutdown():
             _LOGGER.info("Starting graceful shutdown...")

pynumaflow/sink/async_sink.py

@@ -12,6 +12,8 @@
     SINK_SOCK_PATH,
     MAX_MESSAGE_SIZE,
 )
+from pynumaflow.info.server import get_sdk_version, write as info_server_write
+from pynumaflow.info.types import ServerInfo, Protocol, Language, SERVER_INFO_FILE_PATH
 from pynumaflow.sink import Responses, Datum, Response
 from pynumaflow.sink.proto import udsink_pb2_grpc, udsink_pb2
 from pynumaflow.types import NumaflowServicerContext
@@ -124,6 +126,13 @@ def start(self) -> None:
         udsink_pb2_grpc.add_UserDefinedSinkServicer_to_server(Sink(self.__sink_handler), server)
         server.add_insecure_port(self.sock_path)
         server.start()
+        serv_info = ServerInfo(
+            protocol=Protocol.UDS,
+            language=Language.PYTHON,
+            version=get_sdk_version(),
+        )
+        info_server_write(server_info=serv_info, info_file=SERVER_INFO_FILE_PATH)
+
         _LOGGER.info(
             "GRPC Server listening on: %s with max threads: %s", self.sock_path, self._max_threads
         )