From e3eca96d51edcb613f67b48773a711d6486aa1ce Mon Sep 17 00:00:00 2001 From: EazyReal <8047065+EazyReal@users.noreply.github.com> Date: Sat, 18 Jul 2026 01:50:57 +0000 Subject: [PATCH] fix(rollout): make dynamic refills granular Allow positive refill sizes smaller than the target rollout batch, and cancel unfinished non-partial generation tasks before waiting for SGLang to drain. This prevents one rejected prompt group from launching a full replacement wave or keeping the drain loop alive through new agent turns. --- .github/workflows/pr-test.yml | 4 + .github/workflows/pr-test.yml.j2 | 1 + slime/rollout/sglang_rollout.py | 21 ++-- slime/utils/arguments.py | 7 +- tests/test_megatron_argument_validation.py | 18 ++++ tests/test_sglang_rollout_abort.py | 107 +++++++++++++++++++++ 6 files changed, 145 insertions(+), 13 deletions(-) create mode 100644 tests/test_sglang_rollout_abort.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 45606dd532..7a629d8bcb 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -629,6 +629,10 @@ jobs: "num_gpus": 0, "test_file": "test_rollout_validation.py" }, + { + "num_gpus": 0, + "test_file": "test_sglang_rollout_abort.py" + }, { "num_gpus": 0, "test_file": "test_reloadable_process_group_world.py" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index f58b064830..d65217757a 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -82,6 +82,7 @@ {'test_file': 'test_rm_deepscaler.py', 'num_gpus': 0}, {'test_file': 'test_sample.py', 'num_gpus': 0}, {'test_file': 'test_rollout_validation.py', 'num_gpus': 0}, + {'test_file': 'test_sglang_rollout_abort.py', 'num_gpus': 0}, {'test_file': 'test_reloadable_process_group_world.py', 'num_gpus': 0}, {'test_file': 'test_placement_group.py', 'num_gpus': 0}, {'test_file': 'test_external_sglang_engines.py', 'num_gpus': 0}, diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py index 9dcf5fce85..b733b8475a 100644 --- a/slime/rollout/sglang_rollout.py +++ b/slime/rollout/sglang_rollout.py @@ -334,12 +334,17 @@ async def generate_and_rm_group( async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: - aborted_samples = [] - state = GenerateState(args) assert not state.aborted state.aborted = True + if not args.partial_rollout: + for task in state.pendings: + task.cancel() + if state.pendings: + await asyncio.gather(*state.pendings, return_exceptions=True) + state.pendings.clear() + if parse(sglang_router.__version__) <= parse("0.2.1"): response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/list_workers") urls = response["urls"] @@ -349,15 +354,14 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: await abort_servers_until_idle(urls) - # make sure all the pending tasks are finished + if not args.partial_rollout: + return [] + + aborted_samples = [] count = 0 while state.pendings: done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED) - if not args.partial_rollout: - continue - - # for partial rollout, collect the partial samples into the data buffer for task in done: group = task.result() for sample in group: @@ -366,8 +370,7 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: aborted_samples.append(group) count += len(group) - if args.partial_rollout: - logger.info(f"Collected {count} partial samples into the data buffer") + logger.info(f"Collected {count} partial samples into the data buffer") return aborted_samples diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 0a79b99ded..f1e4058e5f 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -1930,10 +1930,9 @@ def slime_validate_args(args): if args.over_sampling_batch_size is None: args.over_sampling_batch_size = args.rollout_batch_size - assert args.over_sampling_batch_size >= args.rollout_batch_size, ( - f"over_sampling_batch_size {args.over_sampling_batch_size} should be greater than or equal to " - f"rollout_batch_size {args.rollout_batch_size}" - ) + assert ( + args.over_sampling_batch_size > 0 + ), f"over_sampling_batch_size {args.over_sampling_batch_size} should be greater than zero" if args.num_epoch is not None: if args.num_rollout is not None: diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py index 6ebb625c80..cc3afb8d8f 100644 --- a/tests/test_megatron_argument_validation.py +++ b/tests/test_megatron_argument_validation.py @@ -305,6 +305,24 @@ def test_slime_validate_args_preserves_zero_rollout_gpus_without_colocate(monkey assert args.offload_rollout is False +@pytest.mark.unit +def test_slime_validate_args_allows_granular_dynamic_sampling_refill(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + args = make_slime_validate_args(rollout_batch_size=64, over_sampling_batch_size=1) + + module.slime_validate_args(args) + + assert args.over_sampling_batch_size == 1 + + +@pytest.mark.unit +def test_slime_validate_args_rejects_empty_dynamic_sampling_refill(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + + with pytest.raises(AssertionError, match="greater than zero"): + module.slime_validate_args(make_slime_validate_args(over_sampling_batch_size=0)) + + @pytest.mark.unit def test_update_weight_delta_requires_disk_transport(monkeypatch): module = load_slime_arguments_module(monkeypatch) diff --git a/tests/test_sglang_rollout_abort.py b/tests/test_sglang_rollout_abort.py new file mode 100644 index 0000000000..776a49b475 --- /dev/null +++ b/tests/test_sglang_rollout_abort.py @@ -0,0 +1,107 @@ +import asyncio +import types + +import pytest + +try: + from plugin_contracts._shared import install_paths, install_stubs +except ImportError: + from tests.plugin_contracts._shared import install_paths, install_stubs + +install_paths() +install_stubs(with_sglang_router=True, with_transformers=True) + +from slime.rollout import sglang_rollout as rollout + +NUM_GPUS = 0 + + +@pytest.mark.unit +def test_abort_cancels_agent_tasks_before_server_drain(monkeypatch) -> None: + async def scenario() -> None: + events = [] + + async def agent_task() -> None: + try: + await asyncio.Future() + finally: + await asyncio.sleep(0) + events.append("agent-cleanup") + + task = asyncio.create_task(agent_task()) + await asyncio.sleep(0) + state = types.SimpleNamespace(aborted=False, pendings={task}) + monkeypatch.setattr(rollout, "GenerateState", lambda _args: state) + + async def get_router_workers(_url): + events.append("list-workers") + return {"urls": ["http://engine"], "workers": [{"url": "http://engine"}]} + + async def drain_servers(urls): + assert urls == ["http://engine"] + assert task.cancelled() + events.append("server-drain") + + monkeypatch.setattr(rollout, "get", get_router_workers) + monkeypatch.setattr(rollout, "abort_servers_until_idle", drain_servers) + + args = types.SimpleNamespace( + partial_rollout=False, + sglang_router_ip="127.0.0.1", + sglang_router_port=30000, + ) + assert await rollout.abort(args, rollout_id=0) == [] + + assert events == ["agent-cleanup", "list-workers", "server-drain"] + assert state.aborted is True + assert state.pendings == set() + + asyncio.run(scenario()) + + +@pytest.mark.unit +def test_abort_preserves_partial_agent_tasks_until_after_server_drain(monkeypatch) -> None: + async def scenario() -> None: + events = [] + server_drained = asyncio.Event() + sample = types.SimpleNamespace(response="partial response", metadata={}) + + async def agent_task(): + await server_drained.wait() + events.append("agent-finished") + return [sample] + + task = asyncio.create_task(agent_task()) + await asyncio.sleep(0) + state = types.SimpleNamespace(aborted=False, pendings={task}) + monkeypatch.setattr(rollout, "GenerateState", lambda _args: state) + + async def get_router_workers(_url): + return {"urls": ["http://engine"], "workers": [{"url": "http://engine"}]} + + async def drain_servers(urls): + assert urls == ["http://engine"] + assert not task.cancelled() + events.append("server-drain") + server_drained.set() + await asyncio.sleep(0) + + monkeypatch.setattr(rollout, "get", get_router_workers) + monkeypatch.setattr(rollout, "abort_servers_until_idle", drain_servers) + + args = types.SimpleNamespace( + partial_rollout=True, + sglang_router_ip="127.0.0.1", + sglang_router_port=30000, + ) + assert await rollout.abort(args, rollout_id=7) == [[sample]] + + assert events == ["server-drain", "agent-finished"] + assert sample.metadata == {"start_rollout_id": 7} + assert state.pendings == set() + + asyncio.run(scenario()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__]))