From b1064613118c230bd47c04eaf7fdf4be002a5b69 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. Make the complete abort lifecycle exception-safe: failures during worker discovery, server drain, or partial collection cancel and await remaining groups before the original failure is re-raised. --- .github/workflows/pr-test.yml | 4 + .github/workflows/pr-test.yml.j2 | 1 + docs/en/examples/glm4-9B.md | 2 +- docs/en/examples/qwen3-4B.md | 2 +- docs/en/get_started/quick_start.md | 4 +- docs/zh/examples/glm4-9B.md | 2 +- docs/zh/examples/qwen3-4B.md | 4 +- docs/zh/get_started/quick_start.md | 4 +- slime/rollout/sglang_rollout.py | 65 ++++--- slime/utils/arguments.py | 7 +- tests/test_megatron_argument_validation.py | 18 ++ tests/test_sglang_rollout_abort.py | 205 +++++++++++++++++++++ 12 files changed, 278 insertions(+), 40 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/docs/en/examples/glm4-9B.md b/docs/en/examples/glm4-9B.md index 4a934ee144..7804fdd804 100644 --- a/docs/en/examples/glm4-9B.md +++ b/docs/en/examples/glm4-9B.md @@ -243,7 +243,7 @@ slime supports more complex sampling schemes, such as the dynamic sampling in [D slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \ ``` -Here, `over_sampling_batch_size` needs to be greater than `rollout_batch_size`. For example: +`over_sampling_batch_size` controls the number of prompt groups requested at a time and must be positive. For example: ```bash --rollout-batch-size 32 \ diff --git a/docs/en/examples/qwen3-4B.md b/docs/en/examples/qwen3-4B.md index 56711878f6..34eaffe020 100644 --- a/docs/en/examples/qwen3-4B.md +++ b/docs/en/examples/qwen3-4B.md @@ -214,7 +214,7 @@ slime supports more complex sampling schemes, such as the dynamic sampling in [D slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \ ``` -Here, `over_sampling_batch_size` needs to be greater than `rollout_batch_size`. For example, you can configure it as: +`over_sampling_batch_size` controls the number of prompt groups requested at a time and must be positive. For example: ```bash --rollout-batch-size 32 \ diff --git a/docs/en/get_started/quick_start.md b/docs/en/get_started/quick_start.md index d0bd3fb780..94f8b50568 100644 --- a/docs/en/get_started/quick_start.md +++ b/docs/en/get_started/quick_start.md @@ -345,7 +345,7 @@ slime supports more complex sampling strategies, such as dynamic sampling used i slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std ``` -Here `over_sampling_batch_size` needs to be greater than `rollout_batch_size`, for example, configured as: +`over_sampling_batch_size` controls the number of prompt groups requested at a time and must be positive. For example: ```bash --rollout-batch-size 32 \ @@ -367,7 +367,7 @@ def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): ) ``` -If the filtering function is very strict, causing a large number of prompt groups to be discarded, the system will monitor the number of pending tasks in `remaining_batch_size`. Once the number of pending tasks drops below the target number (32) due to too many being discarded, the system will automatically trigger a new round of oversampling, requesting `over_sampling_batch_size` (64) new prompts again to repeat the above process. +`remaining_batch_size` tracks accepted and pending prompt groups. Rejected groups decrement it; when it falls below the target (32), the system requests batches of `over_sampling_batch_size` (64) until the accepted and pending total reaches the target again. ### Partial Rollout diff --git a/docs/zh/examples/glm4-9B.md b/docs/zh/examples/glm4-9B.md index 8cb562c1f8..b1503c5765 100644 --- a/docs/zh/examples/glm4-9B.md +++ b/docs/zh/examples/glm4-9B.md @@ -243,7 +243,7 @@ slime 支持了更复杂的 sampling 方案，例如 [DAPO](https://dapo-sia.git slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \ ``` -这里 `over_sampling_batch_size` 需要大于 ``rollout_batch_size`，例如配置为： +`over_sampling_batch_size` 控制每次请求的 prompt 组数，且必须为正数。例如： ```bash --rollout-batch-size 32 \ diff --git a/docs/zh/examples/qwen3-4B.md b/docs/zh/examples/qwen3-4B.md index 2afc70bfbe..fa2f023ec8 100644 --- a/docs/zh/examples/qwen3-4B.md +++ b/docs/zh/examples/qwen3-4B.md @@ -214,7 +214,7 @@ slime 支持了更复杂的 sampling 方案，例如 [DAPO](https://dapo-sia.git slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \ ``` -这里 `over_sampling_batch_size` 需要大于 ``rollout_batch_size`，例如配置为： +`over_sampling_batch_size` 控制每次请求的 prompt 组数，且必须为正数。例如： ```bash --rollout-batch-size 32 \ @@ -298,4 +298,4 @@ ray job submit ... \ `train.py` 和 `train_async.py` 的差别只在于 train loop 的同步逻辑，我们通过 ray 的异步（`.remote`, `ray.get`）实现了这点。 -⚠️ 在异步训练时，sglang 的性能检测日志与训练日志可能会混到一起，不易区分，可以通过 `--sglang-log-level` 来减少 sglang 的日志。 \ No newline at end of file +⚠️ 在异步训练时，sglang 的性能检测日志与训练日志可能会混到一起，不易区分，可以通过 `--sglang-log-level` 来减少 sglang 的日志。 diff --git a/docs/zh/get_started/quick_start.md b/docs/zh/get_started/quick_start.md index 345a5d10be..93904d1eab 100644 --- a/docs/zh/get_started/quick_start.md +++ b/docs/zh/get_started/quick_start.md @@ -347,7 +347,7 @@ slime 支持更复杂的采样策略，例如 [DAPO](https://dapo-sia.github.io/ slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std ``` -这里 `over_sampling_batch_size` 需要大于 `rollout_batch_size`，例如配置为： +`over_sampling_batch_size` 控制每次请求的 prompt 组数，且必须为正数。例如： ```bash --rollout-batch-size 32 \ @@ -369,7 +369,7 @@ def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): ) ``` -如果过滤函数非常严格，导致大量 prompt 组被丢弃，系统会监控 ` remaining_batch_size` 中待处理的任务数量。一旦待处理的任务数因丢弃过多而降至目标数 (32) 以下，系统会自动触发新一轮的过采样，再次请求 `over_sampling_batch_size` (64) 个新的 prompt 重复上述流程。 +`remaining_batch_size` 记录已接收和待处理的 prompt 组数。每丢弃一组，该值就会减一；当它低于目标值 (32) 时，系统会按 `over_sampling_batch_size` (64) 为一批补充请求，直到已接收和待处理的总数重新达到目标值。 ### Partial Rollout diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py index 9dcf5fce85..bcb84d44cf 100644 --- a/slime/rollout/sglang_rollout.py +++ b/slime/rollout/sglang_rollout.py @@ -334,40 +334,51 @@ 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 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"] - else: - response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers") - urls = [worker["url"] for worker in response["workers"]] - - await abort_servers_until_idle(urls) + done = set() + try: + 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"] + else: + response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers") + urls = [worker["url"] for worker in response["workers"]] - # make sure all the pending tasks are finished - count = 0 - while state.pendings: - done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED) + await abort_servers_until_idle(urls) 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: - if sample.response and "start_rollout_id" not in sample.metadata: - sample.metadata["start_rollout_id"] = rollout_id - aborted_samples.append(group) - count += len(group) - - if args.partial_rollout: - logger.info(f"Collected {count} partial samples into the data buffer") + return [] + + aborted_samples = [] + count = 0 + while state.pendings: + done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED) + + for task in done: + group = task.result() + for sample in group: + if sample.response and "start_rollout_id" not in sample.metadata: + sample.metadata["start_rollout_id"] = rollout_id + aborted_samples.append(group) + count += len(group) + except BaseException: + for task in state.pendings: + task.cancel() + await asyncio.gather(*done, *state.pendings, return_exceptions=True) + state.pendings.clear() + raise + + 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..6c2f8c2a80 --- /dev/null +++ b/tests/test_sglang_rollout_abort.py @@ -0,0 +1,205 @@ +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_pending_groups_before_server_drain(monkeypatch) -> None: + async def scenario() -> None: + events = [] + + async def pending_group() -> None: + try: + await asyncio.Future() + finally: + await asyncio.sleep(0) + events.append("pending-cleanup") + + task = asyncio.create_task(pending_group()) + 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 == ["pending-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_groups_until_after_server_drain(monkeypatch) -> None: + async def scenario() -> None: + events = [] + server_drained = asyncio.Event() + sample = types.SimpleNamespace(response="partial response", metadata={}) + + async def pending_group(): + await server_drained.wait() + events.append("pending-finished") + return [sample] + + task = asyncio.create_task(pending_group()) + 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", "pending-finished"] + assert sample.metadata == {"start_rollout_id": 7} + assert state.pendings == set() + + asyncio.run(scenario()) + + +@pytest.mark.unit +def test_abort_cleans_up_partial_groups_when_cancelled_during_drain(monkeypatch) -> None: + async def scenario() -> None: + events = [] + drain_started = asyncio.Event() + + async def pending_group(): + try: + await asyncio.Future() + finally: + events.append("pending-cleanup") + + pending_task = asyncio.create_task(pending_group()) + await asyncio.sleep(0) + state = types.SimpleNamespace(aborted=False, pendings={pending_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"] + events.append("server-drain-started") + drain_started.set() + await asyncio.Future() + + 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, + ) + abort_task = asyncio.create_task(rollout.abort(args, rollout_id=7)) + await drain_started.wait() + abort_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await abort_task + + assert events == ["server-drain-started", "pending-cleanup"] + assert pending_task.cancelled() + assert state.aborted is True + assert state.pendings == set() + + asyncio.run(scenario()) + + +@pytest.mark.unit +def test_abort_cleans_up_partial_groups_when_one_fails(monkeypatch) -> None: + async def scenario() -> None: + events = [] + server_drained = asyncio.Event() + + async def failed_group(): + await server_drained.wait() + raise RuntimeError("generation failed") + + async def pending_group(): + try: + await asyncio.Future() + finally: + events.append("pending-cleanup") + + failed_task = asyncio.create_task(failed_group()) + pending_task = asyncio.create_task(pending_group()) + await asyncio.sleep(0) + state = types.SimpleNamespace(aborted=False, pendings={failed_task, pending_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"] + 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, + ) + with pytest.raises(RuntimeError, match="generation failed"): + await rollout.abort(args, rollout_id=7) + + assert events == ["server-drain", "pending-cleanup"] + assert pending_task.cancelled() + assert state.pendings == set() + + asyncio.run(scenario()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__]))