Merge branch 'main' into clean-shutdown-unarymap

Author: BulkBeing

packages/pynumaflow/pynumaflow/accumulator/_dtypes.py

@@ -17,9 +17,9 @@ class WindowOperation(IntEnum):
     Enumerate the type of Window operation received.
     """
 
-    OPEN = (0,)
-    CLOSE = (1,)
-    APPEND = (2,)
+    OPEN = 0
+    CLOSE = 1
+    APPEND = 2
 
 
 @dataclass(init=False, slots=True)

packages/pynumaflow/pynumaflow/accumulator/servicer/task_manager.py

@@ -14,6 +14,7 @@
     _AccumulatorBuilderClass,
     AccumulatorAsyncCallable,
     WindowOperation,
+    AccumulatorRequest,
 )
 from pynumaflow.proto.accumulator import accumulator_pb2
 from pynumaflow.shared.asynciter import NonBlockingIterator
@@ -93,7 +94,7 @@ async def stream_send_eof(self):
         for unified_key in task_keys:
             await self.tasks[unified_key].iterator.put(STREAM_EOF)
 
-    async def close_task(self, req):
+    async def close_task(self, req: AccumulatorRequest):
         """
         Closes a running accumulator task for a given key.
         Based on the request we compute the unique key, and then
@@ -104,8 +105,9 @@ async def close_task(self, req):
         3. Wait for all the results from the task to be written to the global result  queue
         4. Remove the task from the tracker
         """
-        d = req.payload
-        keys = d.keys
+        # Use keyed_window.keys for task lookup since payload.keys may be empty
+        # (e.g., CLOSE operations don't carry data, so payload.keys is not populated).
+        keys = req.keyed_window.keys
         unified_key = build_unique_key_name(keys)
         curr_task = self.tasks.get(unified_key, None)
 
@@ -120,14 +122,16 @@ async def close_task(self, req):
             # Put the exception in the result queue
             await self.global_result_queue.put(err)
 
-    async def create_task(self, req):
+    async def create_task(self, req: AccumulatorRequest):
         """
         Creates a new accumulator task for the given request.
         Based on the request we compute a unique key, and then
         it creates a new task or appends the request to the existing task.
         """
         d = req.payload
-        keys = d.keys
+        # Use keyed_window.keys for task lookup — the authoritative key identity
+        # for the window, consistent across all operation types (OPEN, APPEND, CLOSE).
+        keys = req.keyed_window.keys
         unified_key = build_unique_key_name(keys)
         curr_task = self.tasks.get(unified_key, None)
 
@@ -138,7 +142,7 @@ async def create_task(self, req):
             # Create a new result queue for the current task
             # We create a new result queue for each task, so that
             # the results of the accumulator operation can be sent to the
-            # the global result queue, which in turn sends the results
+            # global result queue, which in turn sends the results
             # to the client.
             res_queue = NonBlockingIterator()
 
@@ -172,13 +176,14 @@ async def create_task(self, req):
         # Put the request in the iterator
         await curr_task.iterator.put(d)
 
-    async def send_datum_to_task(self, req):
+    async def send_datum_to_task(self, req: AccumulatorRequest):
         """
         Appends the request to the existing window reduce task.
         If the task does not exist, create it.
         """
         d = req.payload
-        keys = d.keys
+        # Use keyed_window.keys for task lookup to match the key used in create_task/close_task.
+        keys = req.keyed_window.keys
         unified_key = build_unique_key_name(keys)
         result = self.tasks.get(unified_key, None)
         if not result:
@@ -215,9 +220,7 @@ async def __invoke_accumulator(
             # Put the exception in the result queue
             await self.global_result_queue.put(err)
 
-    async def process_input_stream(
-        self, request_iterator: AsyncIterable[accumulator_pb2.AccumulatorRequest]
-    ):
+    async def process_input_stream(self, request_iterator: AsyncIterable[AccumulatorRequest]):
         # Start iterating through the request iterator and create tasks
         # based on the operation type received.
         try:
@@ -226,15 +229,15 @@ async def process_input_stream(
                 request_count += 1
                 # check whether the request is an open, append, or close operation
                 match request.operation:
-                    case int(WindowOperation.OPEN):
+                    case WindowOperation.OPEN:
                         # create a new task for the open operation and
                         # put the request in the task iterator
                         await self.create_task(request)
-                    case int(WindowOperation.APPEND):
+                    case WindowOperation.APPEND:
                         # append the task data to the existing task
                         # if the task does not exist, create a new task
                         await self.send_datum_to_task(request)
-                    case int(WindowOperation.CLOSE):
+                    case WindowOperation.CLOSE:
                         # close the current task for req
                         await self.close_task(request)
                     case _:

packages/pynumaflow/pynumaflow/reducestreamer/async_server.py

@@ -1,8 +1,12 @@
+import asyncio
+import contextlib
 import inspect
+import sys
 
 import aiorun
 import grpc
 
+from pynumaflow.info.server import write as info_server_write
 from pynumaflow.info.types import ServerInfo, ContainerType, MINIMUM_NUMAFLOW_VERSION
 from pynumaflow.proto.reducer import reduce_pb2_grpc
 
@@ -15,6 +19,7 @@
     REDUCE_STREAM_SOCK_PATH,
     REDUCE_STREAM_SERVER_INFO_FILE_PATH,
     MAX_NUM_THREADS,
+    NUMAFLOW_GRPC_SHUTDOWN_GRACE_PERIOD_SECONDS,
 )
 
 from pynumaflow.reducestreamer._dtypes import (
@@ -23,7 +28,7 @@
     ReduceStreamer,
 )
 
-from pynumaflow.shared.server import NumaflowServer, check_instance, start_async_server
+from pynumaflow.shared.server import NumaflowServer, check_instance
 
 
 def get_handler(
@@ -156,6 +161,7 @@ def __init__(
         ]
         # Get the servicer instance for the async server
         self.servicer = AsyncReduceStreamServicer(self.reduce_stream_handler)
+        self._error: BaseException | None = None
 
     def start(self):
         """
@@ -166,6 +172,9 @@ def start(self):
             "Starting Async Reduce Stream Server",
         )
         aiorun.run(self.aexec(), use_uvloop=True, shutdown_callback=self.shutdown_callback)
+        if self._error:
+            _LOGGER.critical("Server exiting due to UDF error: %s", self._error)
+            sys.exit(1)
 
     async def aexec(self):
         """
@@ -178,15 +187,42 @@ async def aexec(self):
         # Create a new async server instance and add the servicer to it
         server = grpc.aio.server(options=self._server_options)
         server.add_insecure_port(self.sock_path)
+
+        # The asyncio.Event must be created here (inside aexec) rather than in __init__,
+        # because it must be bound to the running event loop that aiorun creates.
+        shutdown_event = asyncio.Event()
+        self.servicer.set_shutdown_event(shutdown_event)
+
         reduce_pb2_grpc.add_ReduceServicer_to_server(self.servicer, server)
 
         serv_info = ServerInfo.get_default_server_info()
         serv_info.minimum_numaflow_version = MINIMUM_NUMAFLOW_VERSION[ContainerType.Reducestreamer]
-        await start_async_server(
-            server_async=server,
-            sock_path=self.sock_path,
-            max_threads=self.max_threads,
-            cleanup_coroutines=list(),
-            server_info_file=self.server_info_file,
-            server_info=serv_info,
+
+        await server.start()
+        info_server_write(server_info=serv_info, info_file=self.server_info_file)
+
+        _LOGGER.info(
+            "Async GRPC Reduce Stream Server listening on: %s with max threads: %s",
+            self.sock_path,
+            self.max_threads,
         )
+
+        async def _watch_for_shutdown():
+            """Wait for the shutdown event and stop the server with a grace period."""
+            await shutdown_event.wait()
+            _LOGGER.info("Shutdown signal received, stopping server gracefully...")
+            await server.stop(NUMAFLOW_GRPC_SHUTDOWN_GRACE_PERIOD_SECONDS)
+
+        shutdown_task = asyncio.create_task(_watch_for_shutdown())
+        await server.wait_for_termination()
+
+        # Propagate error so start() can exit with a non-zero code
+        self._error = self.servicer._error
+
+        shutdown_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await shutdown_task
+
+        _LOGGER.info("Stopping event loop...")
+        asyncio.get_event_loop().stop()
+        _LOGGER.info("Event loop stopped")

packages/pynumaflow/pynumaflow/reducestreamer/servicer/async_servicer.py

@@ -3,7 +3,7 @@
 
 from google.protobuf import empty_pb2 as _empty_pb2
 
-from pynumaflow._constants import ERR_UDF_EXCEPTION_STRING
+from pynumaflow._constants import ERR_UDF_EXCEPTION_STRING, _LOGGER
 from pynumaflow.proto.reducer import reduce_pb2, reduce_pb2_grpc
 from pynumaflow.reducestreamer._dtypes import (
     Datum,
@@ -12,7 +12,7 @@
     ReduceRequest,
 )
 from pynumaflow.reducestreamer.servicer.task_manager import TaskManager
-from pynumaflow.shared.server import handle_async_error
+from pynumaflow.shared.server import update_context_err
 from pynumaflow.types import NumaflowServicerContext
 
 
@@ -47,6 +47,12 @@ def __init__(
     ):
         # The Reduce handler can be a function or a builder class instance.
         self.__reduce_handler: ReduceStreamAsyncCallable | _ReduceStreamBuilderClass = handler
+        self._shutdown_event: asyncio.Event | None = None
+        self._error: BaseException | None = None
+
+    def set_shutdown_event(self, event: asyncio.Event):
+        """Wire up the shutdown event created by the server's aexec() coroutine."""
+        self._shutdown_event = event
 
     async def ReduceFn(
         self,
@@ -94,20 +100,50 @@ async def ReduceFn(
             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)
+                    err_msg = f"ReduceStreamError, {ERR_UDF_EXCEPTION_STRING}: {repr(msg)}"
+                    _LOGGER.critical(err_msg, exc_info=True)
+                    update_context_err(context, msg, err_msg)
+                    self._error = msg
+                    if self._shutdown_event is not None:
+                        self._shutdown_event.set()
                     return
                 # Send window EOF response or Window result response
                 # back to the client
                 else:
                     yield msg
+        except GeneratorExit:
+            # ReduceFn is an async generator (it yields messages). When Numaflow closes a
+            # window, gRPC calls .aclose() on this generator, throwing GeneratorExit at
+            # the yield point. This is normal stream lifecycle — return cleanly.
+            return
+        except asyncio.CancelledError:
+            # SIGTERM: aiorun cancelled all tasks. Signal the server to stop so
+            # Server.__del__ doesn't try to schedule on a closed event loop.
+            if self._shutdown_event is not None:
+                self._shutdown_event.set()
+            return
         except BaseException as e:
-            await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING)
+            err_msg = f"ReduceStreamError, {ERR_UDF_EXCEPTION_STRING}: {repr(e)}"
+            _LOGGER.critical(err_msg, exc_info=True)
+            update_context_err(context, e, err_msg)
+            self._error = e
+            if self._shutdown_event is not None:
+                self._shutdown_event.set()
             return
         # Wait for the process_input_stream task to finish for a clean exit
         try:
             await producer
+        except asyncio.CancelledError:
+            if self._shutdown_event is not None:
+                self._shutdown_event.set()
+            return
         except BaseException as e:
-            await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING)
+            err_msg = f"ReduceStreamError, {ERR_UDF_EXCEPTION_STRING}: {repr(e)}"
+            _LOGGER.critical(err_msg, exc_info=True)
+            update_context_err(context, e, err_msg)
+            self._error = e
+            if self._shutdown_event is not None:
+                self._shutdown_event.set()
             return
 
     async def IsReady(

packages/pynumaflow/pynumaflow/reducestreamer/servicer/task_manager.py

@@ -195,6 +195,9 @@ async def __invoke_reduce(
             new_instance = self.__reduce_handler.create()
         try:
             _ = await new_instance(keys, request_iterator, output, md)
+        except asyncio.CancelledError:
+            _LOGGER.info("ReduceStream __invoke_reduce cancelled, returning cleanly")
+            return
         # If there is an error in the reduce operation, log and
         # then send the error to the result queue
         except BaseException as err:
@@ -217,6 +220,9 @@ async def process_input_stream(self, request_iterator: AsyncIterable[reduce_pb2.
                         # append the task data to the existing task
                         # if the task does not exist, create a new task
                         await self.send_datum_to_task(request)
+        except asyncio.CancelledError:
+            _LOGGER.info("ReduceStream process_input_stream cancelled, returning cleanly")
+            return
         # If there is an error in the reduce operation, log and
         # then send the error to the result queue
         except BaseException as e:
@@ -261,6 +267,9 @@ async def process_input_stream(self, request_iterator: AsyncIterable[reduce_pb2.
 
             # Once all tasks are completed, senf EOF the global result queue
             await self.global_result_queue.put(STREAM_EOF)
+        except asyncio.CancelledError:
+            _LOGGER.info("ReduceStream post-processing cancelled, returning cleanly")
+            return
         except BaseException as e:
             err_msg = f"Reduce Streaming Error: {repr(e)}"
             _LOGGER.critical(err_msg, exc_info=True)

packages/pynumaflow/tests/accumulator/test_async_accumulator.py

@@ -30,9 +30,11 @@
 def request_generator(count, request, resetkey: bool = False, send_close: bool = False):
     for i in range(count):
         if resetkey:
-            # Clear previous keys and add new ones
+            # Update keys on both payload and keyedWindow to match real platform behavior
             del request.payload.keys[:]
             request.payload.keys.extend([f"key-{i}"])
+            del request.operation.keyedWindow.keys[:]
+            request.operation.keyedWindow.keys.extend([f"key-{i}"])
 
         # Set operation based on index - first is OPEN, rest are APPEND
         if i == 0:
@@ -52,9 +54,11 @@ def request_generator(count, request, resetkey: bool = False, send_close: bool =
 def request_generator_append_only(count, request, resetkey: bool = False):
     for i in range(count):
         if resetkey:
-            # Clear previous keys and add new ones
+            # Update keys on both payload and keyedWindow to match real platform behavior
             del request.payload.keys[:]
             request.payload.keys.extend([f"key-{i}"])
+            del request.operation.keyedWindow.keys[:]
+            request.operation.keyedWindow.keys.extend([f"key-{i}"])
 
         # Set operation to APPEND for all requests
         request.operation.event = accumulator_pb2.AccumulatorRequest.WindowOperation.Event.APPEND
@@ -64,9 +68,11 @@ def request_generator_append_only(count, request, resetkey: bool = False):
 def request_generator_mixed(count, request, resetkey: bool = False):
     for i in range(count):
         if resetkey:
-            # Clear previous keys and add new ones
+            # Update keys on both payload and keyedWindow to match real platform behavior
             del request.payload.keys[:]
             request.payload.keys.extend([f"key-{i}"])
+            del request.operation.keyedWindow.keys[:]
+            request.operation.keyedWindow.keys.extend([f"key-{i}"])
 
         if i % 2 == 0:
             # Set operation to APPEND for even requests
@@ -107,17 +113,26 @@ def start_request() -> accumulator_pb2.AccumulatorRequest:
 
 def start_request_without_open() -> accumulator_pb2.AccumulatorRequest:
     event_time_timestamp, watermark_timestamp = get_time_args()
-
+    window = accumulator_pb2.KeyedWindow(
+        start=mock_interval_window_start(),
+        end=mock_interval_window_end(),
+        slot="slot-0",
+        keys=["test_key"],
+    )
     payload = accumulator_pb2.Payload(
         keys=["test_key"],
         value=mock_message(),
         event_time=event_time_timestamp,
         watermark=watermark_timestamp,
         id="test_id",
     )
-
+    operation = accumulator_pb2.AccumulatorRequest.WindowOperation(
+        event=accumulator_pb2.AccumulatorRequest.WindowOperation.Event.APPEND,
+        keyedWindow=window,
+    )
     request = accumulator_pb2.AccumulatorRequest(
         payload=payload,
+        operation=operation,
     )
     return request