ezmsg 3.7.2__tar.gz → 3.8.0__tar.gz

{ezmsg-3.7.2 → ezmsg-3.8.0}/.github/workflows/ci.yml

@@ -15,7 +15,7 @@ jobs:
   pytest:
     strategy:
       matrix:
-        python-version: ["3.10.15", "3.11"]  # Omit until shutdown fixed: , "3.12", "3.13"]
+        python-version: ["3.10.15", "3.11", "3.12", "3.13"]
         os:
           - "ubuntu-latest"
           - "windows-latest"

{ezmsg-3.7.2 → ezmsg-3.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ezmsg
-Version: 3.7.2
+Version: 3.8.0
 Summary: A simple DAG-based computation model
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>
 License: MIT License

ezmsg-3.8.0/examples/simple_async_publisher.py

@@ -0,0 +1,33 @@
+import asyncio
+
+import ezmsg.core as ez
+
+TOPIC = "/TEST"
+
+
+async def main(host: str = "127.0.0.1", port: int = 12345) -> None:
+    async with ez.GraphContext((host, port), auto_start=True) as ctx:
+        pub = await ctx.publisher(TOPIC)
+        try:
+            print("Publisher Task Launched")
+            count = 0
+            while True:
+                await pub.broadcast(f"{count=}")
+                await asyncio.sleep(0.1)
+                count += 1
+        except asyncio.CancelledError:
+            pass
+        finally:
+            print("Publisher Task Concluded")
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument("--host", default="127.0.0.1", help="hostname for graphserver")
+    parser.add_argument("--port", default=12345, type=int, help="port for graphserver")
+
+    args = parser.parse_args()
+
+    asyncio.run(main(host=args.host, port=args.port))

ezmsg-3.8.0/examples/simple_async_subscriber.py

@@ -0,0 +1,34 @@
+import asyncio
+
+import ezmsg.core as ez
+
+TOPIC = "/TEST"
+
+
+async def main(host: str = "127.0.0.1", port: int = 12345) -> None:
+    async with ez.GraphContext((host, port), auto_start=True) as ctx:
+        sub = await ctx.subscriber(TOPIC)
+        try:
+            print("Subscriber Task Launched")
+            while True:
+                async with sub.recv_zero_copy() as msg:
+                    # Uncomment if you want to witness backpressure!
+                    # await asyncio.sleep(1.0)
+                    print(msg)
+        except asyncio.CancelledError:
+            pass
+        finally:
+            print("Subscriber Task Concluded")
+            print("Detached")
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument("--host", default="127.0.0.1", help="hostname for graphserver")
+    parser.add_argument("--port", default=12345, type=int, help="port for graphserver")
+
+    args = parser.parse_args()
+
+    asyncio.run(main(host=args.host, port=args.port))

ezmsg-3.8.0/examples/simple_publisher.py

@@ -0,0 +1,35 @@
+import time
+
+import ezmsg.core as ez
+
+TOPIC = "/TEST"
+
+def main(host: str = "127.0.0.1", port: int = 12345) -> None:
+    with ez.sync.init((host, port), auto_start=True) as ctx:
+        pub = ctx.create_publisher(TOPIC, force_tcp=True)
+
+        print("Publisher Task Launched")
+        count = 0
+        try:
+            while True:
+                output = f"{count=}"
+                pub.publish(output)
+                print(output)
+                time.sleep(0.1)
+                count += 1
+        except KeyboardInterrupt:
+            pass
+        print("Publisher Task Concluded")
+
+    print("Done")
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument("--host", default="127.0.0.1", help="hostname for graphserver")
+    parser.add_argument("--port", default=12345, type=int, help="port for graphserver")
+    args = parser.parse_args()
+
+    main(host=args.host, port=args.port)

ezmsg-3.8.0/examples/simple_subscriber.py

@@ -0,0 +1,30 @@
+import ezmsg.core as ez
+
+TOPIC = "/TEST"
+
+
+def main(host: str = "127.0.0.1", port: int = 12345) -> None:
+    with ez.sync.init((host, port), auto_start=True) as ctx:
+        print("Subscriber Task Launched")
+
+        def on_message(msg: str) -> None:
+            # Uncomment if you want to witness backpressure!
+            # import time
+            # time.sleep(1.0)
+            print(msg)
+
+        ctx.create_subscription(TOPIC, callback=on_message)
+        ez.sync.spin(ctx)
+
+    print("Subscriber Task Concluded")
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument("--host", default="127.0.0.1", help="hostname for graphserver")
+    parser.add_argument("--port", default=12345, type=int, help="port for graphserver")
+    args = parser.parse_args()
+
+    main(host=args.host, port=args.port)

{ezmsg-3.7.2 → ezmsg-3.8.0}/pyproject.toml

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

{ezmsg-3.7.2 → ezmsg-3.8.0}/src/ezmsg/core/__init__.py

@@ -25,6 +25,10 @@ __all__ = [
     "NormalTermination",
     "GraphServer",
     "GraphContext",
+    "sync",
+    "SyncContext",
+    "SyncPublisher",
+    "SyncSubscriber",
     "run_command",
     "Publisher",
     "Subscriber",
@@ -45,6 +49,8 @@ from .backend import run, GraphRunner, GraphRunnerStartError
 from .backendprocess import Complete, NormalTermination
 from .graphserver import GraphServer
 from .graphcontext import GraphContext
+from . import sync
+from .sync import SyncContext, SyncPublisher, SyncSubscriber
 from .command import run_command
 from .pubclient import Publisher
 from .subclient import Subscriber

{ezmsg-3.7.2 → ezmsg-3.8.0}/src/ezmsg/core/backend.py

@@ -4,6 +4,7 @@ from collections.abc import Collection as AbstractCollection
 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 .graphcontext import GraphContext
 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 @@ class GraphRunner:
         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
@@ -222,6 +226,15 @@ class GraphRunner:
             return self._graph_context.graph_address
         return self._graph_address
 
+    @property
+    def strict_shutdown(self) -> bool:
+        value = os.environ.get("EZMSG_STRICT_SHUTDOWN", "")
+        return value.lower() in ("1", "true", "yes", "on")
+
+    @strict_shutdown.setter
+    def strict_shutdown(self, value: bool) -> None:
+        os.environ["EZMSG_STRICT_SHUTDOWN"] = "1" if value else "0"
+
     @property
     def graph_server_spawned(self) -> bool:
         return self._graph_server_spawned
@@ -311,7 +324,10 @@ class GraphRunner:
         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:
@@ -380,12 +396,15 @@ class GraphRunner:
         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..."
             )
@@ -395,6 +414,8 @@ class GraphRunner:
                 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()
@@ -402,10 +423,72 @@ class GraphRunner:
                     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:

{ezmsg-3.7.2 → ezmsg-3.8.0}/src/ezmsg/core/backendprocess.py

@@ -2,15 +2,19 @@ import asyncio
 import concurrent.futures
 import logging
 import inspect
+import os
 import time
 import traceback
 import threading
+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
-from copy import deepcopy
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures.thread import _worker
 from multiprocessing import Process
 from multiprocessing.synchronize import Event as EventType
 from multiprocessing.synchronize import Barrier as BarrierType
@@ -19,8 +23,7 @@ from concurrent.futures import TimeoutError
 from typing import Any
 
 from .stream import Stream, InputStream, OutputStream
-from .unit import Unit, TIMEIT_ATTR, SUBSCRIBES_ATTR, ZERO_COPY_ATTR
-from .messagechannel import LeakyQueue
+from .unit import Unit, TIMEIT_ATTR, SUBSCRIBES_ATTR
 
 from .graphcontext import GraphContext
 from .pubclient import Publisher
@@ -29,6 +32,75 @@ from .netprotocol import AddressType
 
 logger = logging.getLogger("ezmsg")
 
+STRICT_SHUTDOWN_ENV = "EZMSG_STRICT_SHUTDOWN"
+
+
+def _strict_shutdown_enabled() -> bool:
+    value = os.environ.get(STRICT_SHUTDOWN_ENV, "")
+    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
+        num_threads = len(self._threads)
+        if num_threads >= self._max_workers:
+            return
+        thread_name = f"{self._thread_name_prefix or 'ThreadPool'}_{num_threads}"
+        thread = threading.Thread(
+            name=thread_name,
+            target=_worker,
+            args=(
+                weakref.ref(self),
+                self._work_queue,
+                self._initializer,
+                self._initargs,
+            ),
+        )
+        thread.daemon = True
+        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):
     """
@@ -146,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()
@@ -228,9 +302,14 @@ class DefaultBackendProcess(BackendProcess):
                             ),
                             loop=loop,
                         ).result()
+
+        except asyncio.CancelledError:
+            pass
+
         except Exception:
             self.start_barrier.abort()
-            logger.error(f"{traceback.format_exc()}")
+            # logger.error(f"{traceback.format_exc()}")
+            raise
 
         try:
             logger.debug("Waiting at start barrier!")
@@ -260,9 +339,9 @@ class DefaultBackendProcess(BackendProcess):
                     fn(unit)
                 except NormalTermination:
                     self.term_ev.set()
-                except Exception:
-                    logger.error(f"Exception in Main: {unit.address}")
-                    logger.error(traceback.format_exc())
+                # except Exception:
+                #     logger.error(f"Exception in Main: {unit.address}")
+                #     logger.error(traceback.format_exc())
 
             while True:
                 try:
@@ -309,12 +388,24 @@ class DefaultBackendProcess(BackendProcess):
                     else:
                         unit.shutdown()  # type: ignore
 
-            asyncio.run_coroutine_threadsafe(shutdown_units(), loop=loop).result()
-
-            # for cache in MessageCache.values():
-            #     cache.clear()
+            shutdown_future = asyncio.run_coroutine_threadsafe(shutdown_units(), loop=loop)
+            try:
+                shutdown_future.result()
+            except KeyboardInterrupt:
+                logger.warning("Interrupted during unit shutdown. This may indicate units with slow shutdown methods."
+                               "Re-trying... Press ctrl-c again to terminate immediately.")
+                shutdown_future.result()
 
-            asyncio.run_coroutine_threadsafe(context.revert(), loop=loop).result()
+            revert_future = asyncio.run_coroutine_threadsafe(context.revert(), loop=loop)
+            try:
+                revert_future.result()
+            except KeyboardInterrupt:
+                logger.warning("Interrupted during context revert."
+                               "Re-trying... (will timeout in 10 seconds)")
+                try:
+                    revert_future.result(timeout=10.0)
+                except TimeoutError:
+                    logger.warning("Timed out waiting for retry on context revert")
 
             logger.debug(f"Remaining tasks in event loop = {asyncio.all_tasks(loop)}")
 
@@ -340,6 +431,7 @@ class DefaultBackendProcess(BackendProcess):
         self, unit: Unit, task: Callable
     ) -> Callable[..., Coroutine[Any, Any, None]]:
         task_address = f"{unit.address}:{task.__name__}"
+        strict_shutdown = _strict_shutdown_enabled()
 
         async def publish(stream: Stream, obj: Any) -> None:
             if stream.address in self.pubs:
@@ -374,8 +466,6 @@ class DefaultBackendProcess(BackendProcess):
                 result = call_fn(msg)
                 if inspect.isasyncgen(result):
                     async for stream, obj in result:
-                        if obj and getattr(task, ZERO_COPY_ATTR, False) and obj is msg:
-                            obj = deepcopy(obj)
                         await pub_fn(stream, obj)
 
                 elif asyncio.iscoroutine(result):
@@ -390,9 +480,17 @@ class DefaultBackendProcess(BackendProcess):
                 self.term_ev.set()
                 raise
 
+            except asyncio.CancelledError:
+                # Normal during shutdown; propagate without logging.
+                raise
+
             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
 
         return wrapped_task
 
@@ -468,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.
@@ -477,10 +576,30 @@ 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]
     """
     loop = asyncio.new_event_loop()
+    strict_shutdown = _strict_shutdown_enabled()
+    shutdown_suppress = threading.Event()
+    suppressed_shutdown_errors = {"count": 0}
+    suppressed_lock = threading.Lock()
+    executor = None
+    if not strict_shutdown:
+        executor = _DaemonThreadPoolExecutor(thread_name_prefix="EZMSG")
+        loop.set_default_executor(executor)
+        def _loop_exception_handler(
+            loop_obj: asyncio.AbstractEventLoop, context: dict
+        ) -> None:
+            if shutdown_suppress.is_set():
+                with suppressed_lock:
+                    suppressed_shutdown_errors["count"] += 1
+                return
+            loop_obj.default_exception_handler(context)
+
+        loop.set_exception_handler(_loop_exception_handler)
     thread = threading.Thread(target=run_loop, name="TaskThread", args=(loop,))
     thread.start()
 
@@ -492,6 +611,67 @@ def new_threaded_event_loop(
             logger.debug("Waiting at event...")
             # ev.wait()
         logger.debug("Stopping and closing task thread")
+
+        if not strict_shutdown:
+            shutdown_suppress.set()
+            # Cancel and await remaining tasks before stopping the loop.
+            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 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, 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 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). "
+                        "Shutdown was NOT clean; re-run with EZMSG_STRICT_SHUTDOWN=1 "
+                        "to debug tasks with poor shutdown behavior.",
+                        suppressed_count,
+                        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()