Merge pull request #225 from ezmsg-org/dev

Prepare for release 3.8.0

Author: cboulay

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ezmsg"
-version = "3.7.3"
+version = "3.8.0"
 description = "A simple DAG-based computation model"
 authors = [
   { name = "Griffin Milsap", email = "griffin.milsap@gmail.com" },

src/ezmsg/core/backend.py

@@ -4,6 +4,7 @@
 import enum
 import logging
 import os
+import signal
 from threading import BrokenBarrierError
 from multiprocessing import Event, Barrier
 from multiprocessing.synchronize import Event as EventType
@@ -23,6 +24,7 @@
 from .backendprocess import (
     BackendProcess,
     DefaultBackendProcess,
+    ShutdownSummary,
     new_threaded_event_loop,
 )
 
@@ -170,6 +172,7 @@ class GraphRunner:
     _graph_context: GraphContext | None
     _loop: asyncio.AbstractEventLoop | None
     _loop_cm: object | None
+    _loop_shutdown_summary: ShutdownSummary | None
     _main_process: BackendProcess | None
     _spawned_processes: list[BackendProcess]
     _start_participant: bool
@@ -208,6 +211,7 @@ def __init__(
         self._graph_context = None
         self._loop = None
         self._loop_cm = None
+        self._loop_shutdown_summary = None
         self._main_process = None
         self._spawned_processes = []
         self._start_participant = False
@@ -320,7 +324,10 @@ def _initialize(self, force_single_process: bool, wait_for_ready: bool) -> bool:
         if self._execution_context is None:
             return False
 
-        self._loop_cm = new_threaded_event_loop()
+        self._loop_shutdown_summary = ShutdownSummary()
+        self._loop_cm = new_threaded_event_loop(
+            shutdown_summary=self._loop_shutdown_summary
+        )
         self._loop = self._loop_cm.__enter__()
 
         try:
@@ -389,12 +396,15 @@ def _run_main_process(self) -> None:
         self._main_process = self._execution_context.processes[0]
         self._start_processes(self._execution_context.processes[1:])
 
+        interrupts = 0
+        forced_sigint = False
         try:
             self._main_process.process(self._loop)
             self._join_spawned_processes()
             logger.info("All processes exited normally")
 
         except KeyboardInterrupt:
+            interrupts += 1
             logger.info(
                 "Attempting graceful shutdown, interrupt again to force quit..."
             )
@@ -404,17 +414,81 @@ def _run_main_process(self) -> None:
                 self._join_spawned_processes()
 
             except KeyboardInterrupt:
+                interrupts += 1
+                forced_sigint = True
                 logger.warning("Interrupt intercepted, force quitting")
                 self._execution_context.start_barrier.abort()
                 self._execution_context.stop_barrier.abort()
                 for proc in self._spawned_processes:
                     proc.terminate()
 
         finally:
-            self._join_spawned_processes()
-            self._cleanup()
+            while True:
+                try:
+                    self._join_spawned_processes()
+                    self._cleanup()
+                    break
+                except KeyboardInterrupt:
+                    interrupts += 1
+                    if interrupts >= 2:
+                        forced_sigint = True
+                        logger.warning("Interrupt intercepted, force quitting")
+                        if self._execution_context is not None:
+                            self._execution_context.start_barrier.abort()
+                            self._execution_context.stop_barrier.abort()
+                        for proc in self._spawned_processes:
+                            proc.terminate()
+                        self._cleanup()
+                        break
+                    logger.info(
+                        "Interrupt received during cleanup; attempting graceful shutdown..."
+                    )
+                    if self._execution_context is not None:
+                        self._execution_context.term_ev.set()
             self._started = False
             self._stopped = True
+            if interrupts and not forced_sigint and self._shutdown_was_unclean():
+                forced_sigint = True
+            if forced_sigint:
+                self._exit_with_sigint()
+
+    def _shutdown_was_unclean(self) -> bool:
+        main_shutdown_errors = bool(
+            self._main_process is not None
+            and getattr(self._main_process, "_shutdown_errors", False)
+        )
+        summary = self._loop_shutdown_summary
+        loop_unclean = bool(summary is not None and summary.unclean)
+        return main_shutdown_errors or loop_unclean
+
+    def _exit_with_sigint(self) -> None:
+        code = 0xC000013A if os.name == "nt" else 130
+        if os.name == "nt":
+            try:
+                import ctypes
+            except Exception:
+                os._exit(1)
+            try:
+                ctypes.windll.kernel32.ExitProcess(ctypes.c_uint(code).value)
+            except Exception:
+                os._exit(ctypes.c_int32(code).value)
+            return
+
+        prev_handler = None
+        try:
+            prev_handler = signal.getsignal(signal.SIGINT)
+            signal.signal(signal.SIGINT, signal.SIG_DFL)
+            signal.raise_signal(signal.SIGINT)
+        except Exception:
+            raise SystemExit(code)
+        finally:
+            if prev_handler is not None:
+                try:
+                    signal.signal(signal.SIGINT, prev_handler)
+                except Exception:
+                    pass
+
+        raise SystemExit(code)
 
     def _cleanup(self) -> None:
         if self._cleanup_done:

src/ezmsg/core/backendprocess.py

@@ -9,6 +9,7 @@
 import weakref
 
 from abc import abstractmethod
+from dataclasses import dataclass
 from collections import defaultdict
 from collections.abc import Callable, Coroutine, Generator, Sequence
 from functools import wraps, partial
@@ -39,7 +40,30 @@ def _strict_shutdown_enabled() -> bool:
     return value.lower() in ("1", "true", "yes", "on")
 
 
+@dataclass
+class ShutdownSummary:
+    cancelled_tasks: int = 0
+    executor_active: int = 0
+    pending_tasks: int = 0
+    suppressed_errors: int = 0
+    forced_interrupt: bool = False
+
+    @property
+    def unclean(self) -> bool:
+        return bool(
+            self.executor_active
+            or self.pending_tasks
+            or self.suppressed_errors
+            or self.forced_interrupt
+        )
+
+
 class _DaemonThreadPoolExecutor(ThreadPoolExecutor):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._active_count = 0
+        self._active_lock = threading.Lock()
+
     def _adjust_thread_count(self) -> None:
         if self._broken:
             return
@@ -61,6 +85,22 @@ def _adjust_thread_count(self) -> None:
         thread.start()
         self._threads.add(thread)
 
+    def submit(self, fn, /, *args, **kwargs):
+        fut = super().submit(fn, *args, **kwargs)
+        with self._active_lock:
+            self._active_count += 1
+
+        def _decrement(_):
+            with self._active_lock:
+                self._active_count -= 1
+
+        fut.add_done_callback(_decrement)
+        return fut
+
+    def active_count(self) -> int:
+        with self._active_lock:
+            return self._active_count
+
 
 class Complete(Exception):
     """
@@ -178,11 +218,13 @@ class DefaultBackendProcess(BackendProcess):
     """
 
     pubs: dict[str, Publisher]
+    _shutdown_errors: bool
 
     def process(self, loop: asyncio.AbstractEventLoop) -> None:
         main_func = None
         context = GraphContext(self.graph_address)
         coro_callables: dict[str, Callable[[], Coroutine[Any, Any, None]]] = dict()
+        self._shutdown_errors = False
 
         try:
             self.pubs = dict()
@@ -445,6 +487,8 @@ async def wrapped_task(msg: Any = None) -> None:
             except Exception:
                 logger.error(f"Exception in Task: {task_address}")
                 logger.error(traceback.format_exc())
+                if self.term_ev.is_set():
+                    self._shutdown_errors = True
                 if strict_shutdown:
                     raise
 
@@ -522,6 +566,7 @@ def run_loop(loop: asyncio.AbstractEventLoop):
 @contextmanager
 def new_threaded_event_loop(
     ev: threading.Event | None = None,
+    shutdown_summary: ShutdownSummary | None = None,
 ) -> Generator[asyncio.AbstractEventLoop, None, None]:
     """
     Create a new asyncio event loop running in a separate thread.
@@ -531,6 +576,8 @@ def new_threaded_event_loop(
 
     :param ev: Optional event to signal when the loop is ready.
     :type ev: threading.Event | None
+    :param shutdown_summary: Optional shutdown summary object to populate.
+    :type shutdown_summary: ShutdownSummary | None
     :return: Context manager yielding the event loop.
     :rtype: Generator[asyncio.AbstractEventLoop, None, None]
     """
@@ -539,10 +586,10 @@ def new_threaded_event_loop(
     shutdown_suppress = threading.Event()
     suppressed_shutdown_errors = {"count": 0}
     suppressed_lock = threading.Lock()
+    executor = None
     if not strict_shutdown:
-        loop.set_default_executor(
-            _DaemonThreadPoolExecutor(thread_name_prefix="EZMSG")
-        )
+        executor = _DaemonThreadPoolExecutor(thread_name_prefix="EZMSG")
+        loop.set_default_executor(executor)
         def _loop_exception_handler(
             loop_obj: asyncio.AbstractEventLoop, context: dict
         ) -> None:
@@ -568,36 +615,45 @@ def _loop_exception_handler(
         if not strict_shutdown:
             shutdown_suppress.set()
             # Cancel and await remaining tasks before stopping the loop.
-            async def _cancel_remaining() -> int:
+            async def _cancel_remaining(timeout: float = 1.0) -> tuple[int, int]:
                 tasks = [
                     t
                     for t in asyncio.all_tasks()
                     if t is not asyncio.current_task() and not t.done()
                 ]
                 for t in tasks:
                     t.cancel()
-                if tasks:
-                    await asyncio.wait(tasks)
-                return len(tasks)
+                if not tasks:
+                    return 0, 0
+                _, pending = await asyncio.wait(tasks, timeout=timeout)
+                return len(tasks), len(pending)
 
             cancelled_count = 0
+            pending_count = 0
             forced_interrupt = False
             fut = asyncio.run_coroutine_threadsafe(_cancel_remaining(), loop)
             try:
-                cancelled_count = fut.result()
+                cancelled_count, pending_count = fut.result()
             except KeyboardInterrupt:
                 forced_interrupt = True
                 fut.cancel()
             except Exception:
                 cancelled_count = 0
+                pending_count = 0
 
             suppressed_count = suppressed_shutdown_errors["count"]
-            if cancelled_count or suppressed_count or forced_interrupt:
+            if cancelled_count or suppressed_count or forced_interrupt or pending_count:
                 if forced_interrupt and not cancelled_count and not suppressed_count:
                     logger.warning(
                         "Shutdown interrupted; tasks may still be running. "
                         "Re-run with EZMSG_STRICT_SHUTDOWN=1 to debug tasks with poor shutdown behavior."
                     )
+                elif pending_count:
+                    logger.warning(
+                        "Shutdown timed out waiting for %d task(s). "
+                        "Re-run with EZMSG_STRICT_SHUTDOWN=1 to debug tasks with poor shutdown behavior.",
+                        pending_count,
+                    )
                 else:
                     logger.warning(
                         "Shutdown suppressed %d error(s) and cancelled %d task(s). "
@@ -607,6 +663,15 @@ async def _cancel_remaining() -> int:
                         cancelled_count,
                     )
 
+            if shutdown_summary is not None:
+                shutdown_summary.cancelled_tasks = cancelled_count
+                shutdown_summary.pending_tasks = pending_count
+                shutdown_summary.executor_active = (
+                    executor.active_count() if executor is not None else 0
+                )
+                shutdown_summary.suppressed_errors = suppressed_count
+                shutdown_summary.forced_interrupt = forced_interrupt
+
         loop.call_soon_threadsafe(loop.stop)
         thread.join()
         loop.close()

tests/clean_shutdown_examples_runner.py

@@ -142,8 +142,29 @@ def _listen() -> None:
             + ", ".join(sorted(SYSTEMS))
         )
     system = SYSTEMS[case]()
-    print("READY", flush=True)
-    ez.run(SYSTEM=system)
+    runner = ez.GraphRunner(SYSTEM=system)
+    ready_emitted = threading.Event()
+    done = threading.Event()
+
+    def _emit_ready() -> None:
+        if not ready_emitted.is_set():
+            print("READY", flush=True)
+            ready_emitted.set()
+
+    def _watch_ready() -> None:
+        while not done.is_set():
+            if runner.running:
+                _emit_ready()
+                return
+            time.sleep(0.01)
+        _emit_ready()
+
+    threading.Thread(target=_watch_ready, daemon=True).start()
+    try:
+        runner.run_blocking()
+    finally:
+        done.set()
+        _emit_ready()
 
 
 if __name__ == "__main__":