From af7edb328d00549036d3fd95b2f14d680b04f5a1 Mon Sep 17 00:00:00 2001 From: Salma Date: Thu, 22 Oct 2026 22:23:23 +0530 Subject: [PATCH] fix: cancel started rollouts when trajectory_group_batches exits Closing the iterator, propagating a rollout error, or cancelling the consumer left the rollout tasks it had started running with no owner. Cancel and await them on exit, and retrieve sibling exceptions left unread when the first failure is raised. Closes #1140 --- src/art/batches.py | 53 +++++++------ tests/unit/test_batches_cleanup.py | 116 +++++++++++++++++++++++++++++ 2 files changed, 148 insertions(+), 21 deletions(-) create mode 100644 tests/unit/test_batches_cleanup.py diff --git a/src/art/batches.py b/src/art/batches.py index 21798f4b1..82ad27130 100644 --- a/src/art/batches.py +++ b/src/art/batches.py @@ -23,30 +23,41 @@ async def trajectory_group_batches( ) unstarted = list(groups)[batch_size * skip_batches :] pending = set[asyncio.Task[TrajectoryGroup | None]]() + done = set[asyncio.Task[TrajectoryGroup | None]]() batch = list[TrajectoryGroup]() context = GatherContext( max_exceptions=max_batch_exceptions, increment_pbar=False, ) with set_gather_context(context): - while unstarted or pending or len(batch) >= batch_size: - if context.pbar is None: - context.pbar = tqdm.tqdm(desc=pbar_desc, total=batch_size) - while len(pending) < batch_size * max_concurrent_batches and unstarted: - pending.add(asyncio.create_task(wrap_group_awaitable(unstarted.pop(0)))) - if pending: - done, pending = await asyncio.wait( - pending, - timeout=0 if len(batch) >= batch_size else None, - return_when=asyncio.FIRST_COMPLETED, - ) - context.pbar.update(len(done)) - batch.extend(g for task in done if (g := task.result()) is not None) - if len(batch) >= batch_size: - if context.pbar is not None: - context.pbar.close() - context.reset() - yield batch[:batch_size] - batch = batch[batch_size:] - if batch: - yield batch + try: + while unstarted or pending or len(batch) >= batch_size: + if context.pbar is None: + context.pbar = tqdm.tqdm(desc=pbar_desc, total=batch_size) + while len(pending) < batch_size * max_concurrent_batches and unstarted: + pending.add( + asyncio.create_task(wrap_group_awaitable(unstarted.pop(0))) + ) + if pending: + done, pending = await asyncio.wait( + pending, + timeout=0 if len(batch) >= batch_size else None, + return_when=asyncio.FIRST_COMPLETED, + ) + context.pbar.update(len(done)) + batch.extend(g for task in done if (g := task.result()) is not None) + if len(batch) >= batch_size: + if context.pbar is not None: + context.pbar.close() + context.reset() + yield batch[:batch_size] + batch = batch[batch_size:] + if batch: + yield batch + finally: + # Rollouts started here have no other owner. Stop them when the consumer + # closes the iterator, is cancelled, or receives a rollout error, and + # retrieve sibling exceptions that the error left unread. + for task in pending: + task.cancel() + await asyncio.gather(*pending, *done, return_exceptions=True) diff --git a/tests/unit/test_batches_cleanup.py b/tests/unit/test_batches_cleanup.py new file mode 100644 index 000000000..dd2c900e5 --- /dev/null +++ b/tests/unit/test_batches_cleanup.py @@ -0,0 +1,116 @@ +import asyncio +import gc +from typing import Any + +import pytest + +import art + + +async def completed_group(identity: int) -> art.TrajectoryGroup: + return art.TrajectoryGroup( + [art.Trajectory(reward=float(identity), metadata={"identity": identity})] + ) + + +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 __call__(self) -> art.TrajectoryGroup: + self.started.set() + try: + await self.release.wait() + except asyncio.CancelledError: + self.cancelled = True + raise + self.finished = True + return await completed_group(-1) + + +async def assert_not_resumed(rollout: BlockedRollout) -> None: + # Releasing a leaked rollout would let it finish after its consumer stopped. + rollout.release.set() + for _ in range(5): + await asyncio.sleep(0) + assert rollout.cancelled + assert not rollout.finished + + +async def test_closing_iterator_cancels_started_rollouts() -> None: + blocked = BlockedRollout() + + async def first_group() -> art.TrajectoryGroup: + await blocked.started.wait() + return await completed_group(0) + + batches = art.trajectory_group_batches( + [first_group(), blocked()], + batch_size=1, + max_concurrent_batches=2, + pbar_desc=None, + ) + assert len(await anext(batches)) == 1 + await batches.aclose() + + await assert_not_resumed(blocked) + + +async def test_rollout_failure_cancels_started_siblings() -> None: + blocked = BlockedRollout() + + async def failed_group() -> art.TrajectoryGroup: + await blocked.started.wait() + raise RuntimeError("rollout failed") + + batches = art.trajectory_group_batches( + [failed_group(), blocked()], batch_size=2, pbar_desc=None + ) + with pytest.raises(RuntimeError, match="rollout failed"): + await anext(batches) + + await assert_not_resumed(blocked) + + +async def test_cancelled_consumer_cancels_started_rollouts() -> None: + blocked = BlockedRollout() + batches = art.trajectory_group_batches([blocked()], batch_size=1, pbar_desc=None) + consumer = asyncio.ensure_future(anext(batches)) + await asyncio.wait_for(blocked.started.wait(), timeout=5) + + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + + await assert_not_resumed(blocked) + + +async def test_failure_retrieves_sibling_exceptions() -> None: + unretrieved: list[dict[str, Any]] = [] + loop = asyncio.get_running_loop() + previous_handler = loop.get_exception_handler() + loop.set_exception_handler(lambda _, context: unretrieved.append(context)) + + async def failed_group(message: str) -> art.TrajectoryGroup: + raise RuntimeError(message) + + try: + batches = art.trajectory_group_batches( + [failed_group("first"), failed_group("second")], + batch_size=2, + pbar_desc=None, + ) + with pytest.raises(RuntimeError): + await anext(batches) + del batches + gc.collect() + await asyncio.sleep(0) + finally: + loop.set_exception_handler(previous_handler) + + assert unretrieved == []