From 3a4a4d2f6b143a5b6aa44597b49a0ed801cbaf49 Mon Sep 17 00:00:00 2001 From: Salma Date: Fri, 23 Oct 2026 02:39:50 +0530 Subject: [PATCH] fix: cancel sibling rollouts when a gather fails gather_trajectory_groups and gather_trajectories awaited asyncio.gather, which re-raises the first over-budget failure without cancelling the other rollouts. They kept running against inference, environments and after_each judges with no owner. Cancel and await them before re-raising. Closes #1144 Co-Authored-By: Claude Opus 5.5 --- src/art/gather.py | 39 ++++++++++++-- tests/unit/test_gather_cleanup.py | 85 +++++++++++++++++++++++++++++++ 2 files changed, 119 insertions(+), 5 deletions(-) create mode 100644 tests/unit/test_gather_cleanup.py diff --git a/src/art/gather.py b/src/art/gather.py index 0fe01f861..f756ae2d7 100644 --- a/src/art/gather.py +++ b/src/art/gather.py @@ -3,7 +3,16 @@ 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 @@ -11,6 +20,8 @@ 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: diff --git a/tests/unit/test_gather_cleanup.py b/tests/unit/test_gather_cleanup.py new file mode 100644 index 000000000..7d568641b --- /dev/null +++ b/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]