From a5fd7e2524c7016a0a895f4f8233d8c5f1026cf5 Mon Sep 17 00:00:00 2001 From: Daniel Peng Date: Wed, 14 Oct 2026 14:45:59 +0800 Subject: [PATCH 1/2] fix: drain full trajectory batches after completion bursts --- src/art/batches.py | 2 +- tests/unit/test_batches_completion_bursts.py | 93 ++++++++++++++++++++ 2 files changed, 94 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_batches_completion_bursts.py diff --git a/src/art/batches.py b/src/art/batches.py index 1389a92fc..5a372f81f 100644 --- a/src/art/batches.py +++ b/src/art/batches.py @@ -39,7 +39,7 @@ async def trajectory_group_batches( ) context.pbar.update(len(done)) batch.extend(g for task in done if (g := task.result()) is not None) - if len(batch) >= batch_size: + while len(batch) >= batch_size: if context.pbar is not None: context.pbar.close() context.reset() diff --git a/tests/unit/test_batches_completion_bursts.py b/tests/unit/test_batches_completion_bursts.py new file mode 100644 index 000000000..02b1b512d --- /dev/null +++ b/tests/unit/test_batches_completion_bursts.py @@ -0,0 +1,93 @@ +from collections.abc import Coroutine, Sequence +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})] + ) + + +async def collect_batches( + groups: Sequence[Coroutine[Any, Any, art.TrajectoryGroup]], + *, + max_concurrent_batches: int = 4, + skip_batches: int = 0, + max_batch_exceptions: int = 0, +) -> list[list[art.TrajectoryGroup]]: + try: + return [ + batch + async for batch in art.trajectory_group_batches( + groups, + batch_size=2, + max_concurrent_batches=max_concurrent_batches, + skip_batches=skip_batches, + max_batch_exceptions=max_batch_exceptions, + pbar_desc=None, + ) + ] + finally: + # Skipped inputs are never awaited by the generator and remain caller-owned. + for group in groups: + group.close() + + +def identities(batches: list[list[art.TrajectoryGroup]]) -> list[int]: + return sorted( + group.trajectories[0].metadata["identity"] + for batch in batches + for group in batch + ) + + +@pytest.mark.parametrize( + ("count", "concurrency", "skip", "expected_sizes"), + [ + (0, 4, 0, []), + (1, 4, 0, [1]), + (3, 4, 0, [2, 1]), + (7, 4, 0, [2, 2, 2, 1]), + (8, 4, 0, [2, 2, 2, 2]), + (12, 4, 0, [2, 2, 2, 2, 2, 2]), + (8, 1, 0, [2, 2, 2, 2]), + (8, 4, 1, [2, 2, 2]), + ], +) +async def test_completion_bursts_keep_batches_bounded( + count: int, concurrency: int, skip: int, expected_sizes: list[int] +) -> None: + batches = await collect_batches( + [completed_group(identity) for identity in range(count)], + max_concurrent_batches=concurrency, + skip_batches=skip, + ) + + assert [len(batch) for batch in batches] == expected_sizes + assert identities(batches) == list(range(skip * 2, count)) + + +async def failed_group() -> art.TrajectoryGroup: + raise RuntimeError("rollout failed") + + +async def test_completion_burst_preserves_tolerated_exception() -> None: + # All inputs finish in the same wait, within the original exception budget. + groups = [completed_group(identity) for identity in range(7)] + groups.append(failed_group()) + batches = await collect_batches(groups, max_batch_exceptions=1) + + assert [len(batch) for batch in batches] == [2, 2, 2, 1] + assert identities(batches) == list(range(7)) + + +async def test_completion_burst_propagates_exception_over_budget() -> None: + groups = [completed_group(identity) for identity in range(7)] + groups.append(failed_group()) + + with pytest.raises(RuntimeError, match="rollout failed"): + await collect_batches(groups) From bcf79cd8d22a03b2bc276b414e5fd477ae401de0 Mon Sep 17 00:00:00 2001 From: Daniel Peng Date: Wed, 14 Oct 2026 15:11:39 +0800 Subject: [PATCH 2/2] fix: preserve rollout scheduling between buffered batches --- src/art/batches.py | 17 +++--- tests/unit/test_batches_completion_bursts.py | 63 ++++++++++++++++++++ 2 files changed, 73 insertions(+), 7 deletions(-) diff --git a/src/art/batches.py b/src/art/batches.py index 5a372f81f..21798f4b1 100644 --- a/src/art/batches.py +++ b/src/art/batches.py @@ -29,17 +29,20 @@ async def trajectory_group_batches( increment_pbar=False, ) with set_gather_context(context): - while unstarted or pending: + 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)))) - done, pending = await asyncio.wait( - pending, return_when=asyncio.FIRST_COMPLETED - ) - context.pbar.update(len(done)) - batch.extend(g for task in done if (g := task.result()) is not None) - while len(batch) >= batch_size: + 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() diff --git a/tests/unit/test_batches_completion_bursts.py b/tests/unit/test_batches_completion_bursts.py index 02b1b512d..2a95c27bb 100644 --- a/tests/unit/test_batches_completion_bursts.py +++ b/tests/unit/test_batches_completion_bursts.py @@ -1,3 +1,4 @@ +import asyncio from collections.abc import Coroutine, Sequence from typing import Any @@ -91,3 +92,65 @@ async def test_completion_burst_propagates_exception_over_budget() -> None: with pytest.raises(RuntimeError, match="rollout failed"): await collect_batches(groups) + + +async def test_buffered_batch_checks_completed_pending_failure() -> None: + release_failure = asyncio.Event() + failure_finished = asyncio.Event() + + async def delayed_failure() -> art.TrajectoryGroup: + await release_failure.wait() + failure_finished.set() + raise RuntimeError("pending rollout failed") + + groups = [completed_group(identity) for identity in range(4)] + groups.append(delayed_failure()) + batches = art.trajectory_group_batches(groups, batch_size=2, pbar_desc=None) + try: + assert len(await anext(batches)) == 2 + release_failure.set() + await asyncio.wait_for(failure_finished.wait(), timeout=5) + with pytest.raises(RuntimeError, match="pending rollout failed"): + await anext(batches) + finally: + release_failure.set() + try: + async for _ in batches: + pass + except RuntimeError: + pass + for group in groups: + group.close() + + +async def test_buffered_batch_refills_rollout_slots() -> None: + started = asyncio.Event() + release = asyncio.Event() + + async def next_group(identity: int) -> art.TrajectoryGroup: + started.set() + await release.wait() + return await completed_group(identity) + + groups = [completed_group(identity) for identity in range(8)] + groups.extend(next_group(identity) for identity in range(8, 10)) + batches = art.trajectory_group_batches(groups, batch_size=2, pbar_desc=None) + observed = [] + second = None + try: + observed.append(await anext(batches)) + second = asyncio.ensure_future(anext(batches)) + # A barrier checks scheduling while the consumer handles a buffered batch. + await asyncio.wait_for(started.wait(), timeout=5) + release.set() + observed.append(await second) + observed.extend([batch async for batch in batches]) + assert identities(observed) == list(range(10)) + finally: + release.set() + if second is not None: + await second + async for _ in batches: + pass + for group in groups: + group.close()