Ensure spawned libcoro tasks complete when cancelling

pyproject.toml

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # See file LICENSE for terms.
 
@@ -12,6 +12,9 @@ ignore-words-list = "inout,unparseable,falsy,couldn,Couldn,thirdparty"
 builtin = "clear"
 quiet-level = 3
 
+[tool.cython-lint]
+max-line-length = 120
+
 [tool.isort]
 # Define known first party modules explicitly to avoid directory name dependency.
 # Without this the isort/cython checks may fail if the repository clone is not

python/rapidsmpf/rapidsmpf/streaming/core/CMakeLists.txt

@@ -1,11 +1,11 @@
 # =================================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =================================================================================
 
-set(cython_modules actor.pyx channel.pyx context.pyx fanout.pyx leaf_actor.pyx
+set(cython_modules actor.pyx cancellation.pyx channel.pyx context.pyx fanout.pyx leaf_actor.pyx
                    memory_reserve_or_wait.pyx message.pyx spillable_messages.pyx
 )
 

python/rapidsmpf/rapidsmpf/streaming/core/actor.pyx

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from cpython.object cimport PyObject
@@ -11,7 +11,7 @@ from libcpp.vector cimport vector
 import asyncio
 import inspect
 from collections.abc import Iterable, Mapping
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Future, ThreadPoolExecutor
 from functools import partial, wraps
 
 from rapidsmpf._detail.exception_handling cimport ex_handler
@@ -20,6 +20,8 @@ from rapidsmpf.streaming._detail.libcoro_spawn_task cimport cpp_set_py_future
 from rapidsmpf.streaming.chunks.utils cimport py_deleter
 from rapidsmpf.streaming.core.context cimport Context, cpp_Context
 
+from rapidsmpf.streaming.core.cancellation import (await_cpp_future,
+                                                   shutdown_channels)
 from rapidsmpf.streaming.core.channel import Channel
 from rapidsmpf.streaming.core.context import Context
 
@@ -155,8 +157,7 @@ async def py_actor(func, extra_channels, /, *args, **kwargs):
     try:
         return await func(*args, **kwargs)
     finally:
-        for ch in channels_to_shutdown:
-            await ch.shutdown(ctx)
+        await shutdown_channels(ctx, *channels_to_shutdown)
 
 
 cdef decorate_actor(extra_channels, func):
@@ -224,7 +225,28 @@ def define_actor(*, extra_channels=()):
     return partial(decorate_actor, extra_channels)
 
 
-def sync_wait(coro):
+async def run_and_publish_task(coro, task_ready):
+    """
+    Run a coroutine and expose its event loop and task to another thread.
+
+    Parameters
+    ----------
+    coro
+        Coroutine to run.
+    task_ready
+        Rendezvous point shared with calling thread. Receives the running
+        event loop and cancellable task before ``coro`` starts, allowing
+        the caller to cancel the task from outside the event loop.
+
+    Returns
+    -------
+    Result returned by ``coro``.
+    """
+    task_ready.set_result((asyncio.get_running_loop(), asyncio.current_task()))
+    return await coro
+
+
+def sync_wait(coro, task_ready):
     """
     Run, and wait for completion of, a coroutine.
 
@@ -237,8 +259,13 @@ def sync_wait(coro):
     This should always be called from a thread we control to ensure no live
     event loop is running.
     """
-    with asyncio.Runner() as runner:
-        runner.run(coro)
+    try:
+        with asyncio.Runner() as runner:
+            return runner.run(run_and_publish_task(coro, task_ready))
+    except BaseException as error:
+        if not task_ready.done():
+            task_ready.set_exception(error)
+        raise
 
 
 async def when_all(Context ctx not None, list cpp_actors):
@@ -268,21 +295,7 @@ async def when_all(Context ctx not None, list cpp_actors):
             cpp_set_py_future,
             move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter))
         )
-    try:
-        # This shield makes sure that if a cancellation is raised, the
-        # future is not cancelled.
-        await asyncio.shield(ret)
-    except asyncio.CancelledError as cancel:
-        # The outer awaitable was cancelled, but we must still ensure the
-        # C++ future still runs to completion.
-        try:
-            # This could still fail so we catch and reraise the
-            # cancellation but recording the C++ exception as well.
-            await asyncio.shield(ret)
-        except Exception as cpp_except:
-            raise cancel from cpp_except
-        # Otherwise just reraise the cancellation
-        raise cancel
+    await await_cpp_future(ret)
 
 
 def run_actor_network(Context ctx not None, *, actors):
@@ -344,5 +357,22 @@ def run_actor_network(Context ctx not None, *, actors):
     py_actors = [when_all(ctx, cpp_actors), *py_actors]
     # Need to run in a separate thread in case the cluster runtime already
     # has an async event loop.
+    task_ready = Future()
     with ThreadPoolExecutor(max_workers=1) as executor:
-        executor.submit(sync_wait, run_py_actors(py_actors)).result()
+        worker = executor.submit(sync_wait, run_py_actors(py_actors), task_ready)
+        loop, task = task_ready.result()
+
+        try:
+            worker.result()
+        except BaseException as error:
+            if worker.done():
+                raise
+            # Inner task is not done, let's cancel it and wait for cleanup to complete
+            loop.call_soon_threadsafe(task.cancel)
+            try:
+                worker.result()
+            except asyncio.CancelledError:
+                pass
+            except BaseException as cleanup:
+                raise error from cleanup
+            raise

python/rapidsmpf/rapidsmpf/streaming/core/cancellation.pyi

@@ -0,0 +1,13 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import asyncio
+from collections.abc import Awaitable, Callable
+
+from rapidsmpf.streaming.core.channel import Channel
+from rapidsmpf.streaming.core.context import Context
+
+async def shutdown_channels(ctx: Context, *chs: Channel) -> None: ...
+async def await_cpp_future(
+    future: asyncio.Future[None], *, on_cancel: Callable[[], Awaitable[None]] | None
+) -> None: ...

python/rapidsmpf/rapidsmpf/streaming/core/cancellation.pyx

@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+import asyncio
+
+from rapidsmpf.streaming.core.context cimport Context
+
+
+async def shutdown_channels(Context ctx not None, *chs):
+    """
+    Shutdown channels, recording and then propagating any exceptions
+
+    Parameters
+    ----------
+    ctx
+        Streaming context for channel shutdown
+    chs
+        Channels to shutdown
+
+    Raises
+    ------
+    Any exceptions that shutting down the channels produces.
+    """
+    results = await asyncio.gather(
+        *(ch.shutdown(ctx) for ch in chs), return_exceptions=True
+    )
+    errors = [r for r in results if isinstance(r, BaseException)]
+    if errors:
+        cause = (
+            errors[0]
+            if len(errors) == 1
+            else ExceptionGroup("Errors during shutdown of Channels", errors)
+        )
+        raise cause
+
+
+async def await_cpp_future(future, *, on_cancel=None):
+    """
+    Await a C++ future handling cancellation.
+
+    Parameters
+    ----------
+    future
+        Awaitable future returning None bridging into libcoro.
+    on_cancel
+        Optional callback to run if cancellation is raised while awaiting
+        the future. Must be a zero-argument function that returns an
+        awaitable.
+    """
+    try:
+        # This shield makes sure that if a cancellation is raised, the
+        # future is not cancelled.
+        await asyncio.shield(future)
+    except asyncio.CancelledError as cancelled:
+        # The outer awaitable was cancelled, but we must still ensure the
+        # C++ future still runs to completion.
+        errors = []
+        if on_cancel is not None:
+            try:
+                # Run cancellation callback, e.g. to shutdown channels that are live
+                await on_cancel()
+            except BaseException as error:
+                errors.append(error)
+        try:
+            # This could still fail so we catch and reraise the
+            # cancellation but recording the C++ exception as well.
+            await asyncio.shield(future)
+        except BaseException as error:
+            errors.append(error)
+        if errors:
+            cause = (
+                errors[0]
+                if len(errors) == 1
+                else ExceptionGroup("Errors during cancellation of C++ awaitable", errors)
+            )
+            raise cancelled from cause
+        # Otherwise just reraise the cancellation
+        raise

python/rapidsmpf/rapidsmpf/streaming/core/channel.pyx

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from cpython.object cimport PyObject
@@ -15,6 +15,8 @@ from rapidsmpf.streaming.core.message cimport Message, cpp_Message
 
 import asyncio
 
+from rapidsmpf.streaming.core.cancellation import await_cpp_future
+
 
 cdef extern from * nogil:
     """
@@ -372,7 +374,7 @@ cdef class Channel:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter)),
             )
-        await ret
+        await await_cpp_future(ret, on_cancel=lambda: self.shutdown(ctx))
 
     async def drain_metadata(self, Context ctx not None):
         """
@@ -392,7 +394,7 @@ cdef class Channel:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter)),
             )
-        await ret
+        await await_cpp_future(ret, on_cancel=lambda: self.shutdown_metadata(ctx))
 
     async def shutdown(self, Context ctx not None):
         """
@@ -418,7 +420,7 @@ cdef class Channel:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter)),
             )
-        await ret
+        await await_cpp_future(ret)
 
     async def shutdown_metadata(self, Context ctx not None):
         """
@@ -440,7 +442,7 @@ cdef class Channel:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter)),
             )
-        await ret
+        await await_cpp_future(ret)
 
     async def send(self, Context ctx, Message msg not None):
         """
@@ -467,7 +469,7 @@ cdef class Channel:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter)),
             )
-        await ret
+        await await_cpp_future(ret, on_cancel=lambda: self.shutdown(ctx))
 
     async def send_metadata(self, Context ctx, Message msg not None):
         """
@@ -494,7 +496,7 @@ cdef class Channel:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter)),
             )
-        await ret
+        await await_cpp_future(ret, on_cancel=lambda: self.shutdown_metadata(ctx))
 
     async def recv(self, Context ctx not None):
         """
@@ -520,7 +522,7 @@ cdef class Channel:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter))
             )
-        await ret
+        await await_cpp_future(ret, on_cancel=lambda: self.shutdown(ctx))
         if deref(c_msg).empty():
             return None
         return Message.from_handle(move(deref(c_msg)))
@@ -549,7 +551,7 @@ cdef class Channel:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter))
             )
-        await ret
+        await await_cpp_future(ret, on_cancel=lambda: self.shutdown_metadata(ctx))
         if deref(c_msg).empty():
             return None
         return Message.from_handle(move(deref(c_msg)))

python/rapidsmpf/rapidsmpf/streaming/core/memory_reserve_or_wait.pyx

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from cpython.object cimport PyObject
@@ -22,6 +22,7 @@ from rapidsmpf.streaming.core.context cimport Context, cpp_Context
 import asyncio
 
 import rapidsmpf.utils.string
+from rapidsmpf.streaming.core.cancellation import await_cpp_future
 from rapidsmpf.utils.memory import check_reservation_size
 
 # Sentinel indicating that net_memory_delta estimation has not yet been implemented.
@@ -337,7 +338,7 @@ cdef class MemoryReserveOrWait:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter)),
             )
-        await ret
+        await await_cpp_future(ret)
 
     async def reserve_or_wait(self, size_t size, *, int64_t net_memory_delta):
         """
@@ -407,7 +408,8 @@ cdef class MemoryReserveOrWait:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>future, py_deleter))
             )
-        await future
+        await await_cpp_future(future, on_cancel=self.shutdown)
+
         if not c_ret:
             assert False, "something went wrong, task returned a null pointer!"
         return MemoryReservation.from_handle(move(deref(c_ret)), self._br)
@@ -458,7 +460,8 @@ cdef class MemoryReserveOrWait:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>future, py_deleter)),
             )
-        await future
+        await await_cpp_future(future, on_cancel=self.shutdown)
+
         if not c_ret:
             assert False, "something went wrong, task returned a null pointer!"
         return (
@@ -512,7 +515,8 @@ cdef class MemoryReserveOrWait:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>future, py_deleter))
             )
-        await future
+        await await_cpp_future(future, on_cancel=self.shutdown)
+
         if not c_ret:
             assert False, "something went wrong, task returned a null pointer!"
         return MemoryReservation.from_handle(move(deref(c_ret)), self._br)