diff --git a/docs/api/process/process.md b/docs/api/process/process.md index ac59ad35..abcca2c6 100644 --- a/docs/api/process/process.md +++ b/docs/api/process/process.md @@ -4,3 +4,14 @@ - Process - LocalProcess - RayProcess + +`RayProcess` runs concurrent lifecycle calls in an `asyncio.TaskGroup`. When a +call fails, it cancels sibling waits, requests cancellation of the corresponding +Ray calls, and raises an `ExceptionGroup` containing the failures observed before +cancellation. It does not wait for every component to fail or finish normally. + +Remote cancellation is cooperative: actor methods must yield to their event loop +to be interrupted. Remote cleanup may still be running when the exception is +raised, and local component attributes may reflect an interrupted operation. +Attribute refreshes during error propagation are limited to five seconds; refresh +errors are logged so the original failure is preserved. diff --git a/plugboard/process/ray_process.py b/plugboard/process/ray_process.py index 81a3e28f..bbe835a9 100644 --- a/plugboard/process/ray_process.py +++ b/plugboard/process/ray_process.py @@ -1,6 +1,7 @@ """Provides the `RayProcess` class for managing components in a Ray cluster.""" import asyncio +import sys import typing as _t from plugboard.component import Component @@ -9,7 +10,7 @@ from plugboard.process.process import Process from plugboard.schemas import Resource, Status from plugboard.state import RayStateBackend, StateBackend -from plugboard.utils import build_actor_wrapper, depends_on_optional, gather_except, gen_rand_str +from plugboard.utils import build_actor_wrapper, depends_on_optional, gen_rand_str try: @@ -18,8 +19,38 @@ pass +_ATTRIBUTE_UPDATE_TIMEOUT_SECONDS = 5.0 + + +async def _gather[T](*awaitables: _t.Awaitable[T]) -> list[T]: + """Gather Ray references and local awaitables, cancelling siblings on failure. + + ObjectRefs are awaitable but are not coroutines accepted by TaskGroup. + Cancelling their local waits does not cancel the remote calls, so explicitly + cancel submitted references when the group fails or its caller is cancelled. + """ + + async def _await(awaitable: _t.Awaitable[T]) -> T: + return await awaitable + + try: + async with asyncio.TaskGroup() as group: + tasks = [group.create_task(_await(awaitable)) for awaitable in awaitables] + except BaseException: + for awaitable in awaitables: + if isinstance(awaitable, ray.ObjectRef): + ray.cancel(awaitable) + raise + return [task.result() for task in tasks] + + class RayProcess(Process): - """`RayProcess` manages components in a process model on a multiple Ray actors.""" + """Manages components on multiple Ray actors. + + Concurrent lifecycle calls use TaskGroup: a failure cancels sibling local + waits, requests cancellation of remote calls, and raises an ExceptionGroup. + Remote cancellation is cooperative and may still be in progress on return. + """ _default_state_cls = RayStateBackend @@ -75,12 +106,22 @@ def _create_component_actor(self, component: Component) -> _t.Any: return ray.remote(**ray_options)(actor_cls).remote(**args) # type: ignore - async def _update_component_attributes(self) -> None: - """Updates attributes on local components from remote actors.""" + async def _update_component_attributes(self, *, best_effort: bool = False) -> None: + """Updates local attributes, bounding refreshes while propagating a failure.""" component_ids = [c.id for c in self.components.values()] - remote_states = await gather_except( - *[self._component_actors[id].dict.remote() for id in component_ids] - ) + try: + timeout = _ATTRIBUTE_UPDATE_TIMEOUT_SECONDS if best_effort else None + async with asyncio.timeout(timeout): + remote_states = await _gather( + *[self._component_actors[id].dict.remote() for id in component_ids] + ) + except Exception: + if not best_effort: + raise + self._logger.warning( + "Could not refresh component attributes after failure", exc_info=True + ) + return for id, state in zip(component_ids, remote_states): self.components[id].__dict__.update( { @@ -96,7 +137,7 @@ async def _connect_components(self) -> None: connect_coros = [ component.io_connect.remote(connectors) for component in self._component_actors.values() ] - await gather_except(*connect_coros) + await _gather(*connect_coros) # Allow time for connections to be established # TODO : Replace with a more robust mechanism await asyncio.sleep(1) @@ -109,7 +150,7 @@ async def _connect_state(self) -> None: connector_coros = [ self._state.upsert_connector(connector) for connector in self.connectors.values() ] - await gather_except(*component_coros, *connector_coros) + await _gather(*component_coros, *connector_coros) async def init(self) -> None: """Performs component initialisation actions.""" @@ -117,9 +158,9 @@ async def init(self) -> None: await self._connect_components() coros = [component.init.remote() for component in self._component_actors.values()] try: - await gather_except(*coros) + await _gather(*coros) finally: - await self._update_component_attributes() + await self._update_component_attributes(best_effort=sys.exception() is not None) await super().init() self._logger.info("Process initialised") @@ -128,14 +169,14 @@ async def step(self) -> None: await super().step() coros = [component.step.remote() for component in self._component_actors.values()] try: - await gather_except(*coros) + await _gather(*coros) except Exception: await self._set_status(Status.FAILED) raise else: await self._set_status(Status.WAITING) finally: - await self._update_component_attributes() + await self._update_component_attributes(best_effort=sys.exception() is not None) async def run(self) -> None: """Runs the process to completion.""" @@ -144,10 +185,12 @@ async def run(self) -> None: coros = [component.run.remote() for component in self._component_actors.values()] try: self._tasks = {comp.id: ref for comp, ref in zip(self.components.values(), coros)} - await gather_except(*coros) + await _gather(*coros) except* ray.exceptions.TaskCancelledError: # Ray tasks were cancelled, now call cancel on components to update status - ray.get([component.cancel.remote() for component in self._component_actors.values()]) + await _gather( + *[component.cancel.remote() for component in self._component_actors.values()] + ) except* Exception: await self._set_status(Status.FAILED) raise @@ -156,7 +199,7 @@ async def run(self) -> None: await self._set_status(Status.COMPLETED) finally: self._remove_signal_handlers() - await self._update_component_attributes() + await self._update_component_attributes(best_effort=sys.exception() is not None) self._logger.info("Process run complete") def cancel(self) -> None: @@ -168,5 +211,5 @@ def cancel(self) -> None: async def destroy(self) -> None: """Performs tear-down actions for the `RayProcess` and its `Component`s.""" coros = [component.destroy.remote() for component in self._component_actors.values()] - await gather_except(*coros) + await _gather(*coros) await super().destroy() diff --git a/tests/integration/test_process_with_components_run.py b/tests/integration/test_process_with_components_run.py index 2bb5e0de..d5dbdf32 100644 --- a/tests/integration/test_process_with_components_run.py +++ b/tests/integration/test_process_with_components_run.py @@ -314,14 +314,13 @@ async def test_io_read_with_process_failure( exceptions = exc_info.value.exceptions if process_cls == RayProcess: - # For Ray, we expect both the component failure and the process status error - assert len(exceptions) == 2 + # TaskGroup reports the original failure without waiting for the consumer's + # periodic status check. Concurrent failures may also be included. underlying_errors = [] for e in exceptions: if hasattr(e, "cause") and e.cause: underlying_errors.append(type(e.cause)) assert RuntimeError in underlying_errors - assert ProcessStatusError in underlying_errors else: # For LocalProcess, we only expect the component failure initially assert len(exceptions) == 1 @@ -329,8 +328,7 @@ async def test_io_read_with_process_failure( assert isinstance(inner_exception, RuntimeError) assert "Component failing_comp failed after 2 steps" in str(inner_exception) - # TODO : Change logic of process run to prevent cancellation (similar to Ray)? - # # Consumer should now raise ProcessStatusError when trying to read + # Consumer should now raise ProcessStatusError when trying to read. with pytest.raises( ProcessStatusError, match="Process in failed state for component consumer" ): @@ -339,8 +337,8 @@ async def test_io_read_with_process_failure( # Verify the failing component status is FAILED assert failing_comp.status == Status.FAILED - # Verify consumer status is now STOPPED - assert consumer.status == Status.STOPPED + if process_cls == LocalProcess: + assert consumer.status == Status.STOPPED # The process status should be updated to FAILED due to the failing component process_status = await process.state.get_process_status(process.id) diff --git a/tests/integration/test_ray_process_taskgroup.py b/tests/integration/test_ray_process_taskgroup.py new file mode 100644 index 00000000..6bb04a9a --- /dev/null +++ b/tests/integration/test_ray_process_taskgroup.py @@ -0,0 +1,94 @@ +"""Tests for failure propagation and cancellation of actual Ray actor calls.""" + +import asyncio + +import pytest +import ray + +from plugboard.process.ray_process import _gather + + +@ray.remote +class WaitingActor: + """Actor with a blocked call and observable cancellation.""" + + def __init__(self) -> None: + self.started = asyncio.Event() + self.cancelled = asyncio.Event() + + async def wait(self) -> None: + """Wait indefinitely until cancelled.""" + self.started.set() + try: + await asyncio.Event().wait() + finally: + self.cancelled.set() + + async def wait_started(self) -> None: + """Wait for the blocked call to begin.""" + await self.started.wait() + + async def wait_cancelled(self) -> None: + """Wait for remote cancellation cleanup.""" + await self.cancelled.wait() + + async def fail(self) -> None: + """Fail after the blocked call has started.""" + await self.started.wait() + raise ValueError("actor failed") + + async def echo(self, value: int) -> int: + """Return a value for ordered result checks.""" + return value + + +@pytest.mark.parametrize("remote_failure", [False, True]) +async def test_gather_cancels_remote_sibling(ray_ctx: None, remote_failure: bool) -> None: + """Both local and remote failures must cancel blocked remote work.""" + actor = WaitingActor.remote() + + async def fail_locally() -> None: + await actor.wait_started.remote() + raise ValueError("local failed") + + try: + async with asyncio.timeout(15): + with pytest.raises(ExceptionGroup) as exc_info: + await _gather( + actor.wait.remote(), + actor.fail.remote() if remote_failure else fail_locally(), + ) + assert len(exc_info.value.exceptions) == 1 + assert isinstance(exc_info.value.exceptions[0], ValueError) + await actor.wait_cancelled.remote() + finally: + ray.kill(actor) + + +async def test_gather_caller_cancellation_cancels_remote_work(ray_ctx: None) -> None: + """Cancelling the driver task must cancel the actor call as well.""" + actor = WaitingActor.remote() + try: + async with asyncio.timeout(15): + task = asyncio.create_task(_gather(actor.wait.remote())) + await actor.wait_started.remote() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await actor.wait_cancelled.remote() + finally: + ray.kill(actor) + + +async def test_gather_mixed_results(ray_ctx: None) -> None: + """Gather local coroutines and ObjectRefs together, retaining input order.""" + actor = WaitingActor.remote() + + async def local() -> int: + return 2 + + try: + async with asyncio.timeout(15): + assert await _gather(actor.echo.remote(1), local(), actor.echo.remote(3)) == [1, 2, 3] + finally: + ray.kill(actor) diff --git a/tests/unit/test_ray_process.py b/tests/unit/test_ray_process.py new file mode 100644 index 00000000..a24e38c9 --- /dev/null +++ b/tests/unit/test_ray_process.py @@ -0,0 +1,108 @@ +"""Tests for TaskGroup coordination in RayProcess.""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from plugboard.process import RayProcess +from plugboard.process.ray_process import _gather +from plugboard.schemas import Status + + +async def test_gather_preserves_result_order() -> None: + """Return results in input order even when completion order differs.""" + ready = asyncio.Event() + + async def first() -> int: + await ready.wait() + return 1 + + async def second() -> int: + ready.set() + return 2 + + assert await _gather(first(), second()) == [1, 2] + assert await _gather() == [] + + +async def test_gather_cancels_siblings_on_failure() -> None: + """A failure must propagate without waiting for an indefinitely blocked sibling.""" + started = asyncio.Event() + stopped = asyncio.Event() + + async def blocked() -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + stopped.set() + + async def fail() -> None: + await started.wait() + raise ValueError("component failed") + + async with asyncio.timeout(5): + with pytest.raises(ExceptionGroup) as exc_info: + await _gather(blocked(), fail()) + + assert stopped.is_set() + assert len(exc_info.value.exceptions) == 1 + assert isinstance(exc_info.value.exceptions[0], ValueError) + + +@pytest.mark.parametrize("unresponsive", [False, True]) +async def test_step_preserves_original_failure( + monkeypatch: pytest.MonkeyPatch, unresponsive: bool +) -> None: + """A failed or blocked attribute refresh must not obscure the step failure.""" + process = object.__new__(RayProcess) + component = MagicMock() + component.id = "component" + actor = MagicMock() + actor.step.remote = AsyncMock(side_effect=ValueError("step failed")) + + async def refresh() -> None: + if unresponsive: + await asyncio.Event().wait() + raise RuntimeError("refresh failed") + + actor.dict.remote = refresh + process._component_actors = {component.id: actor} + process.components = {component.id: component} + process._is_initialised = True + process._state_is_connected = False + process._state = MagicMock() + process._logger = MagicMock() + monkeypatch.setattr("plugboard.process.ray_process._ATTRIBUTE_UPDATE_TIMEOUT_SECONDS", 0.01) + + async with asyncio.timeout(5): + with pytest.raises(ExceptionGroup) as exc_info: + await process.step() + + assert process.status == Status.FAILED + assert len(exc_info.value.exceptions) == 1 + assert isinstance(exc_info.value.exceptions[0], ValueError) + assert str(exc_info.value.exceptions[0]) == "step failed" + process._logger.warning.assert_called_once() + + +async def test_gather_propagates_caller_cancellation() -> None: + """Cancelling the caller cleans up local siblings and propagates CancelledError.""" + started = asyncio.Event() + stopped = asyncio.Event() + + async def blocked() -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + stopped.set() + + async with asyncio.timeout(5): + task = asyncio.create_task(_gather(blocked())) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert stopped.is_set()