fix: fail fast in RayProcess with TaskGroup

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.

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,17 +150,17 @@ 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."""
         await self.connect_state()
         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()

tests/integration/test_process_with_components_run.py

@@ -314,23 +314,21 @@ 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
         inner_exception = exceptions[0]
         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)

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)