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 == []