[WIP]feat: add multiproc async

examples/map/async_multiproc_map/Dockerfile

@@ -0,0 +1,55 @@
+####################################################################################################
+# builder: install needed dependencies
+####################################################################################################
+
+FROM python:3.10-slim-bullseye AS builder
+
+ENV PYTHONFAULTHANDLER=1 \
+  PYTHONUNBUFFERED=1 \
+  PYTHONHASHSEED=random \
+  PIP_NO_CACHE_DIR=on \
+  PIP_DISABLE_PIP_VERSION_CHECK=on \
+  PIP_DEFAULT_TIMEOUT=100 \
+  POETRY_VERSION=1.2.2 \
+  POETRY_HOME="/opt/poetry" \
+  POETRY_VIRTUALENVS_IN_PROJECT=true \
+  POETRY_NO_INTERACTION=1 \
+  PYSETUP_PATH="/opt/pysetup"
+
+ENV EXAMPLE_PATH="$PYSETUP_PATH/examples/map/async_multiproc_map"
+ENV VENV_PATH="$EXAMPLE_PATH/.venv"
+ENV PATH="$POETRY_HOME/bin:$VENV_PATH/bin:$PATH"
+
+RUN apt-get update \
+    && apt-get install --no-install-recommends -y \
+        curl \
+        wget \
+        # deps for building python deps
+        build-essential \
+    && apt-get install -y git \
+    && apt-get clean && rm -rf /var/lib/apt/lists/* \
+    \
+    # install dumb-init
+    && wget -O /dumb-init https://github.com/Yelp/dumb-init/releases/download/v1.2.5/dumb-init_1.2.5_x86_64 \
+    && chmod +x /dumb-init \
+    && curl -sSL https://install.python-poetry.org | python3 -
+
+####################################################################################################
+# udf: used for running the udf vertices
+####################################################################################################
+FROM builder AS udf
+
+WORKDIR $PYSETUP_PATH
+COPY ./ ./
+
+WORKDIR $EXAMPLE_PATH
+RUN poetry lock
+RUN poetry install --no-cache --no-root && \
+    rm -rf ~/.cache/pypoetry/
+
+RUN chmod +x entry.sh
+
+ENTRYPOINT ["/dumb-init", "--"]
+CMD ["sh", "-c", "$EXAMPLE_PATH/entry.sh"]
+
+EXPOSE 5000

examples/map/async_multiproc_map/Makefile

@@ -0,0 +1,22 @@
+TAG ?= v6
+PUSH ?= false
+IMAGE_REGISTRY = quay.io/skohli/numaflow-python/async-multiproc:${TAG}
+DOCKER_FILE_PATH = examples/map/async_multiproc_map/Dockerfile
+
+.PHONY: update
+update:
+	poetry update -vv
+
+.PHONY: image-push
+image-push: update
+	cd ../../../ && docker buildx build \
+	-f ${DOCKER_FILE_PATH} \
+	-t ${IMAGE_REGISTRY} \
+	 --platform linux/amd64,linux/arm64 . --push
+
+.PHONY: image
+image: update
+	cd ../../../ && docker build \
+	-f ${DOCKER_FILE_PATH} \
+	-t ${IMAGE_REGISTRY} .
+	@if [ "$(PUSH)" = "true" ]; then docker push ${IMAGE_REGISTRY}; fi

examples/map/async_multiproc_map/README.md

@@ -0,0 +1,20 @@
+# Multiprocessing Map
+
+`pynumaflow` supports only asyncio based Reduce UDFs because we found that procedural Python is not able to handle 
+any substantial traffic. 
+
+This features enables the `pynumaflow` developer to utilise multiprocessing capabilities while 
+writing UDFs using the map function. These are particularly useful for CPU intensive operations,
+as it allows for better resource utilisation.
+
+In this mode we would spawn N number (N = Cpu count) of grpc servers in different processes, where each of them are
+listening on multiple TCP sockets.
+
+To enable multiprocessing mode start the multiproc server in the UDF using the following command,
+providing the optional argument `server_count` to specify the number of
+servers to be forked (defaults to `os.cpu_count` if not provided):
+```python
+if __name__ == "__main__":
+    grpc_server = MapMultiProcServer(handler, server_count = 3)
+    grpc_server.start()
+```

examples/map/async_multiproc_map/entry.sh

@@ -0,0 +1,4 @@
+#!/bin/sh
+set -eux
+
+python example.py

examples/map/async_multiproc_map/example.py

@@ -0,0 +1,40 @@
+import os
+
+from pynumaflow.mapper import Messages, Message, Datum, Mapper, AsyncMapMultiprocServer
+from pynumaflow._constants import _LOGGER
+
+
+class FlatMap(Mapper):
+    """
+    This class needs to be of type Mapper class to be used
+    as a handler for the MapServer class.
+    Example of a mapper that calculates if a number is prime.
+    """
+
+    async def handler(self, keys: list[str], datum: Datum) -> Messages:
+        val = datum.value
+        _ = datum.event_time
+        _ = datum.watermark
+        messages = Messages()
+        messages.append(Message(val, keys=keys))
+        _LOGGER.info(f"MY PID {os.getpid()}")
+        return messages
+
+
+if __name__ == "__main__":
+    """
+    Example of starting a multiprocessing map vertex.
+    """
+    # To set the env server_count value set the env variable
+    # NUM_CPU_MULTIPROC="N"
+    server_count = int(os.getenv("NUM_CPU_MULTIPROC", "2"))
+    server_type = os.getenv("SERVER_KIND", "tcp")
+    use_tcp = False
+    if server_type == "tcp":
+        use_tcp = True
+    elif server_type == "uds":
+        use_tcp = False
+    _class = FlatMap()
+    # Server count is the number of server processes to start
+    grpc_server = AsyncMapMultiprocServer(_class, server_count=server_count, use_tcp=use_tcp)
+    grpc_server.start()

examples/map/async_multiproc_map/pipeline.yaml

@@ -0,0 +1,42 @@
+apiVersion: numaflow.numaproj.io/v1alpha1
+kind: Pipeline
+metadata:
+  name: simple-pipeline
+spec:
+  limits:
+    readBatchSize: 10
+  vertices:
+    - name: in
+      source:
+        # A self data generating source
+        generator:
+          rpu: 200
+          duration: 1s
+    - name: mult
+      udf:
+        container:
+          image: quay.io/skohli/numaflow-python/async-multiproc:v5
+#          imagePullPolicy: Always
+          env:
+            - name: SERVER_KIND
+              value: "uds"
+            - name: PYTHONDEBUG
+              value: "true"
+            - name: NUM_CPU_MULTIPROC
+              value: "3"  # DO NOT forget the double quotes!!!
+      containerTemplate:
+        env:
+          - name: NUMAFLOW_RUNTIME
+            value: "rust"
+          - name: NUMAFLOW_DEBUG
+            value: "true" # DO NOT forget the double quotes!!!
+
+    - name: out
+      sink:
+        # A simple log printing sink
+        log: {}
+  edges:
+    - from: in
+      to: mult
+    - from: mult
+      to: out

examples/map/async_multiproc_map/pyproject.toml

@@ -0,0 +1,15 @@
+[tool.poetry]
+name = "async-multiproc-forward-message"
+version = "0.2.4"
+description = ""
+authors = ["Numaflow developers"]
+
+[tool.poetry.dependencies]
+python = ">=3.10,<3.13"
+pynumaflow = { path = "../../../"}
+
+[tool.poetry.dev-dependencies]
+
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"

pynumaflow/batchmapper/servicer/async_servicer.py

@@ -98,7 +98,7 @@ async def MapFn(
 
         except BaseException as err:
             _LOGGER.critical("UDFError, re-raising the error", exc_info=True)
-            await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING)
+            await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING, False)
             return
 
     async def IsReady(

pynumaflow/info/types.py

@@ -26,6 +26,7 @@
 # MULTIPROC_KEY is the field used to indicate that Multiproc map mode is enabled
 # The value contains the number of servers spawned.
 MULTIPROC_KEY = "MULTIPROC"
+MULTIPROC_ENDPOINTS = "MULTIPROC_ENDPOINTS"
 
 SI = TypeVar("SI", bound="ServerInfo")
 

pynumaflow/mapper/__init__.py

@@ -1,3 +1,4 @@
+from pynumaflow.mapper.async_multiproc_server import AsyncMapMultiprocServer
 from pynumaflow.mapper.async_server import MapAsyncServer
 from pynumaflow.mapper.multiproc_server import MapMultiprocServer
 from pynumaflow.mapper.sync_server import MapServer
@@ -13,4 +14,5 @@
     "MapServer",
     "MapAsyncServer",
     "MapMultiprocServer",
+    "AsyncMapMultiprocServer",
 ]

pynumaflow/mapper/_servicer/_async_servicer.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 from collections.abc import AsyncIterable
 
 from google.protobuf import empty_pb2 as _empty_pb2
@@ -18,11 +19,10 @@
     Provides the functionality for the required rpc methods.
     """
 
-    def __init__(
-        self,
-        handler: MapAsyncCallable,
-    ):
+    def __init__(self, handler: MapAsyncCallable, multiproc: bool = False):
         self.background_tasks = set()
+        # This indicates whether the grpc server attached is multiproc or not
+        self.multiproc = multiproc
         self.__map_handler: MapAsyncCallable = handler
 
     async def MapFn(
@@ -36,6 +36,7 @@
         """
         # proto repeated field(keys) is of type google._upb._message.RepeatedScalarContainer
         # we need to explicitly convert it to list
+        producer = None
         try:
             # The first message to be received should be a valid handshake
             req = await request_iterator.__anext__()
@@ -56,44 +57,65 @@
             async for msg in consumer:
                 # If the message is an exception, we raise the exception
                 if isinstance(msg, BaseException):
-                    await handle_async_error(context, msg, ERR_UDF_EXCEPTION_STRING)
+                    await handle_async_error(context, msg, ERR_UDF_EXCEPTION_STRING, self.multiproc)
                     return
                 # Send window response back to the client
                 else:
                     yield msg
             # wait for the producer task to complete
             await producer
+        except GeneratorExit:
+            _LOGGER.info("Client disconnected, generator closed.")
+            raise
         except BaseException as e:
             _LOGGER.critical("UDFError, re-raising the error", exc_info=True)
-            await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING)
+            await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING, self.multiproc)
             return
+        finally:
+            if producer and not producer.done():
+                producer.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await producer
 
-    async def _process_inputs(
-        self,
-        request_iterator: AsyncIterable[map_pb2.MapRequest],
-        result_queue: NonBlockingIterator,
-    ):
-        """
-        Utility function for processing incoming MapRequests
-        """
+    async def _process_inputs(self, request_iterator, result_queue):
         try:
-            # for each incoming request, create a background task to execute the
-            # UDF code
             async for req in request_iterator:
-                msg_task = asyncio.create_task(self._invoke_map(req, result_queue))
-                # save a reference to a set to store active tasks
-                self.background_tasks.add(msg_task)
-                msg_task.add_done_callback(self.background_tasks.discard)
-
-            # wait for all tasks to complete
-            for task in self.background_tasks:
-                await task
+                task = asyncio.create_task(self._invoke_map(req, result_queue))
+                self.background_tasks.add(task)
+                task.add_done_callback(self.background_tasks.discard)
 
-            # send an EOF to result queue to indicate that all tasks have completed
+            await asyncio.gather(*self.background_tasks)
+        except BaseException:
+            _LOGGER.critical("MapFn Error in _process_inputs", exc_info=True)
+        finally:
             await result_queue.put(STREAM_EOF)
 
-        except BaseException:
-            _LOGGER.critical("MapFn Error, re-raising the error", exc_info=True)
+    # async def _process_inputs(
+    #     self,
+    #     request_iterator: AsyncIterable[map_pb2.MapRequest],
+    #     result_queue: NonBlockingIterator,
+    # ):
+    #     """
+    #     Utility function for processing incoming MapRequests
+    #     """
+    #     try:
+    #         # for each incoming request, create a background task to execute the
+    #         # UDF code
+    #         async for req in request_iterator:
+    #             msg_task = asyncio.create_task(self._invoke_map(req, result_queue))
+    #             # save a reference to a set to store active tasks
+    #             self.background_tasks.add(msg_task)
+    #             msg_task.add_done_callback(self.background_tasks.discard)
+    #
+    #         # wait for all tasks to complete
+    #         for task in self.background_tasks:
+    #             await task
+    #
+    #         # send an EOF to result queue to indicate that all tasks have completed
+    #         await result_queue.put(STREAM_EOF)
+    #
+    #     except BaseException:
+    #         _LOGGER.critical("MapFn Error, re-raising the error", exc_info=True)
 
     async def _invoke_map(self, req: map_pb2.MapRequest, result_queue: NonBlockingIterator):
         """