fix: update asyncio_reduce example to the current reducer API (Fixes #338)

.github/workflows/build-push.yaml

@@ -18,12 +18,10 @@ jobs:
 
     strategy:
       matrix:
-        # TODO: rewrite asyncio_reduce example using latest SDK version, as it is currently using old methods
-        # then add to example_directories matrix
         example_directories: [
           "examples/map/even_odd", "examples/map/flatmap", "examples/map/forward_message",
           "examples/map/multiproc_map", "examples/mapstream/flatmap_stream", "examples/reduce/counter",
-          "examples/reducestream/counter", "examples/reducestream/sum",
+          "examples/reduce/asyncio_reduce", "examples/reducestream/counter", "examples/reducestream/sum",
           "examples/sideinput/simple_sideinput", "examples/sideinput/simple_sideinput/udf",
           "examples/sink/async_log", "examples/sink/log", "examples/sink/all_sinks",
           "examples/source/simple_source", "examples/sourcetransform/event_time_filter",

packages/pynumaflow/examples/reduce/asyncio_reduce/Dockerfile

@@ -33,6 +33,6 @@ COPY --from=builder $EXAMPLE_PATH/ $EXAMPLE_PATH/
 # is not present (by design, to keep the image small), so uv run fails.
 # Instead, we activate the pre-built .venv via PATH and run python directly.
 ENV PATH="$EXAMPLE_PATH/.venv/bin:$PATH"
-CMD ["python", "example.py"]
+CMD ["python", "async_exec.py"]
 
 EXPOSE 5000

packages/pynumaflow/examples/reduce/asyncio_reduce/README.md

@@ -85,7 +85,8 @@ Default executor pool would work just as well in the use case.
 ```python
 if __name__ == "__main__":
     ...
-    threadPool = ExecutorPool(exec_type=e_type, max_workers=max_workers)
+    exec_pool = ExecutorPool(exec_type=e_type, max_workers=mx_workers)
+    grpc_server = ReduceAsyncServer(ReduceHandler, init_args=(exec_pool,))
     ...
 ```
 

packages/pynumaflow/examples/reduce/asyncio_reduce/async_exec.py

@@ -1,5 +1,4 @@
 import aiohttp
-import aiorun
 import asyncio
 import json
 import time
@@ -11,7 +10,7 @@
     Message,
     Datum,
     Metadata,
-    AsyncReducer,
+    ReduceAsyncServer,
 )
 
 _LOGGER = setup_logging(__name__)
@@ -47,5 +46,5 @@ async def reduce_handler(keys: list[str], datums: AsyncIterable[Datum], md: Meta
 
 
 if __name__ == "__main__":
-    grpc_server = AsyncReducer(handler=reduce_handler)
-    aiorun.run(grpc_server.start())
+    grpc_server = ReduceAsyncServer(reduce_handler)
+    grpc_server.start()

packages/pynumaflow/examples/reduce/asyncio_reduce/executorpool.py

@@ -1,4 +1,3 @@
-import aiorun
 import asyncio
 import concurrent.futures
 import os
@@ -12,33 +11,31 @@
     Message,
     Datum,
     Metadata,
-    AsyncReducer,
+    ReduceAsyncServer,
+    Reducer,
 )
 
 _LOGGER = setup_logging(__name__)
 
 
-class ReduceHandler:
-    def __init__(self, exec_pool=None):
+class ReduceHandler(Reducer):
+    def __init__(self, exec_pool):
         self.exec_pool = exec_pool
 
-    async def reduce_handler(
-        self, key: list[str], datums: AsyncIterable[Datum], md: Metadata
-    ) -> Messages:
+    async def handler(self, key: list[str], datums: AsyncIterable[Datum], md: Metadata) -> Messages:
         """
         handler function for executing operations on the messages received by the reduce vertex
         Here the function is used to execute certain blocking operations to demonstrate the use of
         asyncio with ThreadPool and ProcessPool Executor
         """
         interval_window = md.interval_window
         tasks = []
-        self.exec_pool.set_loop(asyncio.get_event_loop())
         start_time = time.time()
         async for _ in datums:
             url = "http://host.docker.internal:9888/ping"
-            fut = threadPool.submit(blocking_call, url)
+            fut = self.exec_pool.submit(blocking_call, url)
             tasks.append(fut)
-        await threadPool.gather(tasks)
+        await self.exec_pool.gather(tasks)
         end_time = time.time()
         msg = (
             f"batch_time:{end_time - start_time} interval_window_start:{interval_window.start} "
@@ -60,9 +57,6 @@ class ExecutorPool:
         thread - Used to create a ThreadPool Executor
         process - Used to create a ProcessPool Executor
 
-    loop : asyncio event loop
-        the event loop on which the tasks should be submitted
-
     max_workers : Number of workers in the executor pool
         the max number of threads/process to be created in the pool
 
@@ -74,22 +68,18 @@ class ExecutorPool:
     async gather(tasks, return_exceptions, return_when)
         Wait for the coroutines to execute and gather the results
 
-    set_loop(event_loop)
-        Update the event loop associated with the executor
-
     close(wait)
          Shutdown the executor
     """
 
-    def __init__(self, exec_type="thread", max_workers=None, loop=None):
-        self.loop = loop or asyncio.get_event_loop()
+    def __init__(self, exec_type="thread", max_workers=None):
         if exec_type == "thread":
             self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
         elif exec_type == "process":
             self.executor = concurrent.futures.ProcessPoolExecutor(max_workers=max_workers)
 
     def submit(self, func, *args):
-        return self.loop.run_in_executor(self.executor, func, args)
+        return asyncio.get_running_loop().run_in_executor(self.executor, func, args)
 
     async def gather(self, tasks, return_exceptions=False, return_when="ALL_COMPLETED"):
         completed, pending = await asyncio.wait(tasks, return_when=return_when)
@@ -99,9 +89,6 @@ async def gather(self, tasks, return_exceptions=False, return_when="ALL_COMPLETE
         # self.tasks = []
         return results
 
-    def set_loop(self, event_loop):
-        self.loop = event_loop
-
     def close(self, wait=True):
         self.executor.shutdown(wait)
 
@@ -126,15 +113,11 @@ def blocking_call(url):
 EXEC_TYPE is environment var to decide the type of executor to be used by the UDF
 """
 if __name__ == "__main__":
-    e_type = os.getenv("EXEC_TYPE")
-    if e_type:
-        e_type = e_type.lower()
+    e_type = os.getenv("EXEC_TYPE", "thread").lower()
     mx_workers = os.getenv("MAX_WORKERS")
     if mx_workers:
         mx_workers = int(mx_workers)
 
-    threadPool = ExecutorPool(exec_type=e_type, max_workers=mx_workers)
-    handler = ReduceHandler(exec_pool=threadPool)
-
-    grpc_server = AsyncReducer(handler=handler.reduce_handler)
-    aiorun.run(grpc_server.start())
+    exec_pool = ExecutorPool(exec_type=e_type, max_workers=mx_workers)
+    grpc_server = ReduceAsyncServer(ReduceHandler, init_args=(exec_pool,))
+    grpc_server.start()

packages/pynumaflow/examples/reduce/asyncio_reduce/pyproject.toml

@@ -1,12 +1,12 @@
 [project]
-name = "reduce-counter"
+name = "asyncio-reduce"
 version = "0.2.4"
 description = ""
 requires-python = ">=3.13"
 dependencies = [
     "pynumaflow",
-    "aiorun>=2025.1.1",
     "aiohttp>=3.8.4",
+    "requests>=2.32.0",
 ]
 
 [tool.uv.sources]