fix: cancel sibling rollouts when a gather fails

src/art/gather.py

@@ -3,14 +3,25 @@
 import contextlib
 import contextvars
 from dataclasses import dataclass, field
-from typing import Awaitable, Callable, Iterable, Iterator, Literal, Sequence, overload
+from typing import (
+    Awaitable,
+    Callable,
+    Iterable,
+    Iterator,
+    Literal,
+    Sequence,
+    TypeVar,
+    overload,
+)
 
 from openai.types.chat.chat_completion import Choice
 from tqdm import auto as tqdm
 
 from .preprocessing.vllm_tokens import choice_completion_tokens
 from .trajectories import Trajectory, TrajectoryGroup
 
+T = TypeVar("T")
+
 
 async def gather_trajectory_groups(
     groups: Iterable[Awaitable[TrajectoryGroup]],
@@ -45,10 +56,10 @@ async def group_forward(g: Awaitable[TrajectoryGroup]):
         return await after_each(group)
 
     with set_gather_context(context):
-        future = asyncio.gather(*[group_forward(g) for g in groups])
+        tasks = [asyncio.ensure_future(group_forward(g)) for g in groups]
         total = sum(getattr(g, "_num_trajectories", 1) for g in groups)
         context.pbar = tqdm.tqdm(desc=pbar_desc, total=total)
-        result_groups = await future
+        result_groups = await _gather_or_cancel(tasks)
 
     if context.pbar is not None:
         context.pbar.close()
@@ -125,14 +136,32 @@ async def gather_trajectories(
         max_exceptions=max_exceptions,
     )
     with set_gather_context(context):
-        results = await asyncio.gather(
-            *[wrap_trajectories_awaitable(t) for t in trajectories_list]
+        results = await _gather_or_cancel(
+            [
+                asyncio.ensure_future(wrap_trajectories_awaitable(t))
+                for t in trajectories_list
+            ]
         )
     if context.pbar is not None:
         context.pbar.close()
     return results
 
 
+async def _gather_or_cancel(tasks: list[asyncio.Future[T]]) -> list[T]:
+    """Like `asyncio.gather`, but a failure also stops the remaining tasks.
+
+    `asyncio.gather` re-raises the first exception and leaves the other rollouts
+    running with no owner. Cancel and await them before re-raising.
+    """
+    try:
+        return await asyncio.gather(*tasks)
+    except BaseException:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        raise
+
+
 async def wrap_group_awaitable(
     awaitable: Awaitable[TrajectoryGroup],
 ) -> TrajectoryGroup | None:

tests/unit/test_gather_cleanup.py

@@ -0,0 +1,85 @@
+import asyncio
+
+import pytest
+
+import art
+
+
+class BlockedRollout:
+    """A rollout that stays in flight until released and records how it ended."""
+
+    def __init__(self) -> None:
+        self.started = asyncio.Event()
+        self.release = asyncio.Event()
+        self.cancelled = False
+        self.finished = False
+
+    async def wait(self) -> None:
+        self.started.set()
+        try:
+            await self.release.wait()
+        except asyncio.CancelledError:
+            self.cancelled = True
+            raise
+        self.finished = True
+
+
+async def assert_not_resumed(rollout: BlockedRollout) -> None:
+    # Releasing a leaked rollout would let it finish after gather raised.
+    rollout.release.set()
+    for _ in range(5):
+        await asyncio.sleep(0)
+    assert rollout.cancelled
+    assert not rollout.finished
+
+
+async def test_group_failure_cancels_sibling_rollouts() -> None:
+    blocked = BlockedRollout()
+
+    async def slow_group() -> art.TrajectoryGroup:
+        await blocked.wait()
+        return art.TrajectoryGroup([art.Trajectory(reward=1)])
+
+    async def failed_group() -> art.TrajectoryGroup:
+        await blocked.started.wait()
+        raise RuntimeError("rollout failed")
+
+    with pytest.raises(RuntimeError, match="rollout failed"):
+        await art.gather_trajectory_groups(
+            [failed_group(), slow_group()], pbar_desc=None
+        )
+
+    await assert_not_resumed(blocked)
+
+
+async def test_trajectory_failure_cancels_sibling_rollouts() -> None:
+    blocked = BlockedRollout()
+
+    async def slow_trajectory() -> art.Trajectory:
+        await blocked.wait()
+        return art.Trajectory(reward=1)
+
+    async def failed_trajectory() -> art.Trajectory:
+        await blocked.started.wait()
+        raise RuntimeError("rollout failed")
+
+    with pytest.raises(RuntimeError, match="rollout failed"):
+        await art.gather_trajectories(
+            [failed_trajectory(), slow_trajectory()], pbar_desc=None
+        )
+
+    await assert_not_resumed(blocked)
+
+
+async def test_tolerated_failure_still_gathers_every_group() -> None:
+    async def group(reward: float) -> art.TrajectoryGroup:
+        return art.TrajectoryGroup([art.Trajectory(reward=reward)])
+
+    async def failed_group() -> art.TrajectoryGroup:
+        raise RuntimeError("rollout failed")
+
+    groups = await art.gather_trajectory_groups(
+        [group(0.0), failed_group(), group(1.0)], pbar_desc=None, max_exceptions=1
+    )
+
+    assert [g.trajectories[0].reward for g in groups] == [0.0, 1.0]