fix(rollout): make dynamic refills granular

.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"

.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},

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
 

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:

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)

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__]))