multiproc map clean shutdown

Signed-off-by: Sreekanth <prsreekanth920@gmail.com>

Author: BulkBeing

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

@@ -213,6 +213,9 @@ async def __invoke_accumulator(
             _ = await new_instance(request_iterator, output)
             # send EOF to the output stream
             await output.put(STREAM_EOF)
+        except asyncio.CancelledError:
+            # Task cancelled during shutdown (e.g. SIGTERM) — not a UDF fault.
+            return
         # If there is an error in the accumulator operation, log and
         # then send the error to the result queue
         except BaseException as err:
@@ -243,6 +246,9 @@ async def process_input_stream(self, request_iterator: AsyncIterable[Accumulator
                     case _:
                         _LOGGER.debug(f"No operation matched for request: {request}", exc_info=True)
 
+        except asyncio.CancelledError:
+            # Task cancelled during shutdown (e.g. SIGTERM) — not a UDF fault.
+            return
         # If there is an error in the accumulator operation, log and
         # then send the error to the result queue
         except BaseException as e:
@@ -274,6 +280,9 @@ async def process_input_stream(self, request_iterator: AsyncIterable[Accumulator
 
             # Now send STREAM_EOF to terminate the global result queue iterator
             await self.global_result_queue.put(STREAM_EOF)
+        except asyncio.CancelledError:
+            # Task cancelled during shutdown (e.g. SIGTERM) — not a UDF fault.
+            return
         except BaseException as e:
             err_msg = f"Accumulator Streaming Error: {repr(e)}"
             _LOGGER.critical(err_msg, exc_info=True)

packages/pynumaflow/pynumaflow/reducer/servicer/task_manager.py

@@ -165,6 +165,9 @@ async def __invoke_reduce(
             new_instance = self.__reduce_handler.create()
         try:
             msgs = await new_instance(keys, request_iterator, md)
+        except asyncio.CancelledError:
+            # Task cancelled during shutdown (e.g. SIGTERM) — not a UDF fault.
+            raise
         except BaseException as err:
             _LOGGER.critical("UDFError, re-raising the error", exc_info=True)
             err_msg = f"ReduceError: {repr(err)}"

packages/pynumaflow/pynumaflow/shared/server.py

@@ -3,6 +3,7 @@
 import multiprocessing
 import multiprocessing.synchronize
 import os
+import signal
 import socket
 import threading
 import traceback
@@ -191,6 +192,17 @@ def start_multiproc_server(
     server_info.metadata[MULTIPROC_KEY] = str(process_count)
     info_server_write(server_info=server_info, info_file=server_info_file)
 
+    # Register a SIGTERM handler so that kubectl delete triggers graceful
+    # shutdown of all child workers via the shared multiprocessing.Event,
+    # instead of the default abrupt kill.
+    if shutdown_event is not None:
+
+        def _sigterm_handler(signum, frame):
+            _LOGGER.info("SIGTERM received, signalling workers to shut down...")
+            shutdown_event.set()
+
+        signal.signal(signal.SIGTERM, _sigterm_handler)
+
     for worker in workers:
         worker.join()