change reduce signature

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

Author: kohlisid

README.md

@@ -45,7 +45,7 @@ pre-commit install
 - [Implement User Defined Source Transformers](https://github.com/numaproj/numaflow-python/tree/main/examples/sourcetransform)
 - Implement User Defined Functions
     - [Map](https://github.com/numaproj/numaflow-python/tree/main/examples/map)
-    - [Reduce](https://github.com/numaproj/numaflow-python/tree/main/examples/reduce/counter)
+    - [Reduce](https://github.com/numaproj/numaflow-python/tree/main/examples/reduce)
     - [Map Stream](https://github.com/numaproj/numaflow-python/tree/main/examples/mapstream)
 - [Implement User Defined Sinks](https://github.com/numaproj/numaflow-python/tree/main/examples/sink)
 - [Implement User Defined SideInputs](https://github.com/numaproj/numaflow-python/tree/main/examples/sideinput)

examples/reduce/README.md

@@ -0,0 +1,65 @@
+# Reducer in Python
+
+For creating a reducer UDF we can use two different approaches:
+- Class based reducer
+  - For the class based reducer we need to implement a class that inherits from the `Reducer` class and implements the required methods.
+  - Next we need to create a `ReduceAsyncServer` instance and pass the reducer class to it along with any input args or 
+  kwargs that the custom reducer class requires.
+  - Finally we need to call the `start` method on the `ReduceAsyncServer` instance to start the reducer server.
+  ```python
+    from numaflow import Reducer, ReduceAsyncServer
+    class Example(Reducer):
+    def __init__(self, counter):
+        self.counter = counter
+
+    async def handler(
+        self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata
+    ) -> Messages:
+        interval_window = md.interval_window
+        self.counter = 0
+        async for _ in datums:
+            self.counter += 1
+        msg = (
+            f"counter:{self.counter} interval_window_start:{interval_window.start} "
+            f"interval_window_end:{interval_window.end}"
+        )
+        return Messages(Message(str.encode(msg), keys=keys))
+
+    if __name__ == "__main__":
+    # Here we are using the class instance as the reducer_instance
+    # which will be used to invoke the handler function.
+    # We are passing the init_args for the class instance.
+    grpc_server = ReduceAsyncServer(Example, init_args=(0,))
+    grpc_server.start()
+  ``` 
+
+- Function based reducer
+    For the function based reducer we need to create a function of the signature
+    ```python
+    async def handler(keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> Messages:
+    ```
+    that takes the required arguments and returns the `Messages` object.
+    - Next we need to create a `ReduceAsyncServer` instance and pass the function to it along with any input args or kwargs that the custom reducer function requires.
+    - Finally we need to call the `start` method on the `ReduceAsyncServer` instance to start the reducer server.
+    - We must ensure that no init_args or init_kwargs are passed to the `ReduceAsyncServer` instance as they are not used for function based reducers.
+    ```python
+        from numaflow import ReduceAsyncServer
+        async def handler(keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> Messages:
+            counter = 0
+            interval_window = md.interval_window
+            async for _ in datums:
+                counter += 1
+            msg = (
+                f"counter:{counter} interval_window_start:{interval_window.start} "
+                f"interval_window_end:{interval_window.end}"
+            )
+            return Messages(Message(str.encode(msg), keys=keys))
+    
+        if __name__ == "__main__":
+        # Here we are using the function as the reducer_instance
+        # which will be used to invoke the handler function.
+        grpc_server = ReduceAsyncServer(handler)
+        grpc_server.start()
+    ```
+
+

examples/reduce/counter/example.py

@@ -5,8 +5,8 @@
 
 
 class Example(Reducer):
-    def __init__(self):
-        self.counter = 0
+    def __init__(self, counter):
+        self.counter = counter
 
     async def handler(
         self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata
@@ -37,7 +37,11 @@ async def reduce_handler(keys: list[str], datums: AsyncIterable[Datum], md: Meta
 if __name__ == "__main__":
     invoke = os.getenv("INVOKE", "handler")
     if invoke == "class":
-        grpc_server = ReduceAsyncServer(Example())
+        # Here we are using the class instance as the reducer_instance
+        # which will be used to invoke the handler function.
+        # We are passing the init_args for the class instance.
+        grpc_server = ReduceAsyncServer(Example, init_args=(0,))
     else:
+        # Here we are using the handler function directly as the reducer_instance.
         grpc_server = ReduceAsyncServer(reduce_handler)
     grpc_server.start()

pynumaflow/reducer/_dtypes.py

@@ -1,7 +1,6 @@
 from abc import ABCMeta, abstractmethod
 from asyncio import Task
 from collections.abc import Iterator, Sequence, Awaitable
-from copy import deepcopy
 from dataclasses import dataclass
 from datetime import datetime
 from typing import TypeVar, Callable, Union
@@ -234,6 +233,9 @@ def keys(self) -> list[str]:
         return self._key
 
 
+ReduceAsyncCallable = Callable[[list[str], AsyncIterable[Datum], Metadata], Awaitable[Messages]]
+
+
 class Reducer(metaclass=ABCMeta):
     """
     Provides an interface to write a Reducer
@@ -247,17 +249,6 @@ def __call__(self, *args, **kwargs):
         """
         return self.handler(*args, **kwargs)
 
-    def __deepcopy__(self, memo):
-        """
-        Allow to deepcopy the class instance.
-        """
-        cls = self.__class__
-        result = cls.__new__(cls)
-        memo[id(self)] = result
-        for k, v in self.__dict__.items():
-            setattr(result, k, deepcopy(v, memo))
-        return result
-
     @abstractmethod
     async def handler(
         self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata
@@ -268,6 +259,26 @@ async def handler(
         pass
 
 
-ReduceAsyncCallable = Callable[[list[str], AsyncIterable[Datum], Metadata], Awaitable[Messages]]
+class ReduceBuilderClass:
+    """
+    Class to build a Reducer class instance.
+    Args:
+        reducer_class: the reducer class to be used for Reduce UDF
+        args: the arguments to be passed to the reducer class
+        kwargs: the keyword arguments to be passed to the reducer class
+    """
+
+    def __init__(self, reducer_class: type[Reducer], args: tuple, kwargs: dict):
+        self._reducer_class: type[Reducer] = reducer_class
+        self._args = args
+        self._kwargs = kwargs
+
+    def create(self) -> Reducer:
+        """
+        Create a new Reducer instance.
+        """
+        return self._reducer_class(*self._args, **self._kwargs)
+
+
 # ReduceCallable is a callable which can be used as a handler for the Reduce UDF.
-ReduceCallable = Union[ReduceAsyncCallable, Reducer]
+ReduceCallable = Union[ReduceAsyncCallable, type[Reducer]]

pynumaflow/reducer/async_server.py

@@ -1,3 +1,5 @@
+import inspect
+
 import aiorun
 import grpc
 
@@ -12,19 +14,44 @@
     _LOGGER,
 )
 
-from pynumaflow.reducer._dtypes import ReduceCallable
+from pynumaflow.reducer._dtypes import (
+    ReduceCallable,
+    ReduceBuilderClass,
+    Reducer,
+)
 
 from pynumaflow.shared.server import NumaflowServer, start_async_server
 
 
+def get_handler(reducer_handler: ReduceCallable, init_args: tuple = (), init_kwargs: dict = None):
+    """
+    Get the correct handler type based on the arguments passed
+    """
+    if inspect.isfunction(reducer_handler):
+        if len(init_args) > 0 or len(init_kwargs) > 0:
+            # if the init_args or init_kwargs are passed, then the reducer_handler
+            # can only be of class Reducer type
+            raise TypeError("Cannot pass function handler with init args or kwargs")
+        # return the function handler
+        return reducer_handler
+    elif issubclass(reducer_handler, Reducer):
+        # if handler is type of Class Reducer, create a new instance of
+        # a ReducerBuilderClass
+        return ReduceBuilderClass(reducer_handler, init_args, init_kwargs)
+    else:
+        raise TypeError("Invalid type passed")
+
+
 class ReduceAsyncServer(NumaflowServer):
     """
     Class for a new Reduce Server instance.
     """
 
     def __init__(
         self,
-        reducer_instance: ReduceCallable,
+        reducer_handler: ReduceCallable,
+        init_args: tuple = (),
+        init_kwargs: dict = None,
         sock_path=REDUCE_SOCK_PATH,
         max_message_size=MAX_MESSAGE_SIZE,
         max_threads=MAX_THREADS,
@@ -41,7 +68,9 @@ def __init__(
                         defaults to number of processors x4
         server_type: The type of server to be used
         """
-        self.reducer_instance: ReduceCallable = reducer_instance
+        if init_kwargs is None:
+            init_kwargs = {}
+        self.reducer_handler = get_handler(reducer_handler, init_args, init_kwargs)
         self.sock_path = f"unix://{sock_path}"
         self.max_message_size = max_message_size
         self.max_threads = max_threads
@@ -51,7 +80,7 @@ def __init__(
             ("grpc.max_receive_message_length", self.max_message_size),
         ]
         # Get the servicer instance for the async server
-        self.servicer = AsyncReduceServicer(reducer_instance)
+        self.servicer = AsyncReduceServicer(self.reducer_handler)
 
     def start(self):
         """

pynumaflow/reducer/servicer/async_servicer.py

@@ -1,8 +1,8 @@
 import asyncio
-from copy import deepcopy
 
 from datetime import datetime, timezone
 from collections.abc import AsyncIterable
+from typing import Union
 
 import grpc
 from google.protobuf import empty_pb2 as _empty_pb2
@@ -13,8 +13,14 @@
     STREAM_EOF,
     DELIMITER,
 )
-from pynumaflow.reducer._dtypes import Datum, IntervalWindow, Metadata, Reducer
-from pynumaflow.reducer._dtypes import ReduceResult, ReduceCallable
+from pynumaflow.reducer._dtypes import (
+    Datum,
+    IntervalWindow,
+    Metadata,
+    ReduceAsyncCallable,
+    ReduceBuilderClass,
+)
+from pynumaflow.reducer._dtypes import ReduceResult
 from pynumaflow.reducer.servicer.asynciter import NonBlockingIterator
 from pynumaflow.proto.reducer import reduce_pb2, reduce_pb2_grpc
 from pynumaflow.types import NumaflowServicerContext
@@ -43,13 +49,14 @@ class AsyncReduceServicer(reduce_pb2_grpc.ReduceServicer):
 
     def __init__(
         self,
-        handler: ReduceCallable,
+        handler: Union[ReduceAsyncCallable, ReduceBuilderClass],
     ):
         # Collection for storing strong references to all running tasks.
         # Event loop only keeps a weak reference, which can cause it to
         # get lost during execution.
         self.background_tasks = set()
-        self.__reduce_handler: ReduceCallable = handler
+        # The reduce handler can be a function or a builder class instance.
+        self.__reduce_handler: Union[ReduceAsyncCallable, ReduceBuilderClass] = handler
 
     async def ReduceFn(
         self,
@@ -143,11 +150,12 @@ async def __invoke_reduce(
         self, keys: list[str], request_iterator: AsyncIterable[Datum], md: Metadata
     ):
         new_instance = self.__reduce_handler
-        # If the reduce handler is a class instance, create a new copy of it.
+        # If the reduce handler is a class instance, create a new instance of it.
         # It is required for a new key to be processed by a
         # new instance of the reducer for a given window
-        if isinstance(self.__reduce_handler, Reducer):
-            new_instance = deepcopy(self.__reduce_handler)
+        # Otherwise the function handler can be called directly
+        if isinstance(self.__reduce_handler, ReduceBuilderClass):
+            new_instance = self.__reduce_handler.create()
         try:
             msgs = await new_instance(keys, request_iterator, md)
         except Exception as err:

tests/reduce/test_async_reduce.py

@@ -28,24 +28,6 @@
 
 LOGGER = setup_logging(__name__)
 
-# if set to true, map handler will raise a `ValueError` exception.
-raise_error_from_map = False
-
-
-async def async_reduce_handler(
-    keys: list[str], datums: AsyncIterable[Datum], md: Metadata
-) -> Messages:
-    interval_window = md.interval_window
-    counter = 0
-    async for _ in datums:
-        counter += 1
-    msg = (
-        f"counter:{counter} interval_window_start:{interval_window.start} "
-        f"interval_window_end:{interval_window.end}"
-    )
-
-    return Messages(Message(str.encode(msg), keys=keys))
-
 
 def request_generator(count, request, resetkey: bool = False):
     for i in range(count):
@@ -80,8 +62,8 @@ def startup_callable(loop):
 
 
 class ExampleClass(Reducer):
-    def __init__(self):
-        self.counter = 0
+    def __init__(self, counter):
+        self.counter = counter
 
     async def handler(
         self, keys: list[str], datums: AsyncIterable[Datum], md: Metadata
@@ -97,10 +79,20 @@ async def handler(
         return Messages(Message(str.encode(msg), keys=keys))
 
 
-def NewAsyncReducer(
-    reduce_handler=async_reduce_handler,
-):
-    server_instance = ReduceAsyncServer(reducer_instance=ExampleClass())
+async def err_handler(keys: list[str], datums: AsyncIterable[Datum], md: Metadata) -> Messages:
+    interval_window = md.interval_window
+    counter = 0
+    async for _ in datums:
+        counter += 1
+    msg = (
+        f"counter:{counter} interval_window_start:{interval_window.start} "
+        f"interval_window_end:{interval_window.end}"
+    )
+    return Messages(Message(str.encode(msg), keys=keys))
+
+
+def NewAsyncReducer():
+    server_instance = ReduceAsyncServer(ExampleClass, init_args=(0,))
     udfs = server_instance.servicer
 
     return udfs
@@ -240,8 +232,13 @@ def __stub(self):
         return reduce_pb2_grpc.ReduceStub(_channel)
 
     def test_error_init(self):
+        # Check that reducer_handler in required
         with self.assertRaises(TypeError):
             ReduceAsyncServer()
+        # Check that the init_args and init_kwargs are passed
+        # only with a Reducer class
+        with self.assertRaises(TypeError):
+            ReduceAsyncServer(err_handler, init_args=(0, 1))
 
 
 if __name__ == "__main__":

tests/reduce/test_async_reduce_err.py

@@ -73,7 +73,7 @@ async def err_handler(keys: list[str], datums: AsyncIterable[Datum], md: Metadat
 
 
 def NewAsyncReducer():
-    server_instance = ReduceAsyncServer(reducer_instance=err_handler)
+    server_instance = ReduceAsyncServer(err_handler)
     udfs = server_instance.servicer
 
     return udfs