diff --git a/pyproject.toml b/pyproject.toml index 8d22591b5..9cfc8ae6c 100644 --- a/pyproject.toml +++ b/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 diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/CMakeLists.txt b/python/rapidsmpf/rapidsmpf/streaming/core/CMakeLists.txt index 36b3405ae..a62cba81f 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/core/CMakeLists.txt +++ b/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 ) diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/actor.pyx b/python/rapidsmpf/rapidsmpf/streaming/core/actor.pyx index f97aea2c8..8a849e8cf 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/core/actor.pyx +++ b/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(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 diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/cancellation.pyi b/python/rapidsmpf/rapidsmpf/streaming/core/cancellation.pyi new file mode 100644 index 000000000..5f3a367d0 --- /dev/null +++ b/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: ... diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/cancellation.pyx b/python/rapidsmpf/rapidsmpf/streaming/core/cancellation.pyx new file mode 100644 index 000000000..a4d8afc0e --- /dev/null +++ b/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 diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/channel.pyx b/python/rapidsmpf/rapidsmpf/streaming/core/channel.pyx index 17fdf1db6..4b7cba64f 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/core/channel.pyx +++ b/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(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(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(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(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(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(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(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(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))) diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/memory_reserve_or_wait.pyx b/python/rapidsmpf/rapidsmpf/streaming/core/memory_reserve_or_wait.pyx index 21048e406..99b757006 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/core/memory_reserve_or_wait.pyx +++ b/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(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(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(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(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)