rollout: support group-scoped session affinity

Author: chengcuiping

docs/en/advanced/sglang-config.md

@@ -302,20 +302,27 @@ You can configure the routing policy:
 
 ### Session-Affinity Routing for Multi-Turn Agents
 
-For multi-turn dialogues and agentic workloads, session affinity ensures that all requests belonging to the same conversation are routed to the same backend worker. This significantly improves prefix cache hit rates because the worker already has the conversation history cached.
+For multi-turn dialogues and agentic workloads, session affinity routes requests that share a routing identity to the same backend worker. This is a placement control; it does not guarantee a cache hit or a performance improvement.
 
-slime automatically assigns each sample a unique `session_id` (stored in `sample.session_id`). When the router policy is `consistent_hashing`, this ID is passed as the `X-SMG-Routing-Key` header, and SGLang Model Gateway uses it to deterministically route all turns of the same session to the same worker.
+By default, slime automatically assigns each sample a unique `session_id` (stored in `sample.session_id`). Users can instead choose group scope so sibling samples in one rollout group share a routing key. This can be useful for fan-out and grouped multi-turn rollouts. Callers may always set `Sample.session_id` directly; an explicit non-empty ID is preserved.
 
 ```bash
 --router-policy consistent_hashing
+# Default: one generated session ID per sample.
+--rollout-session-id-scope sample
+
+# Opt in: one generated or inherited session ID per rollout group.
+--rollout-session-id-scope group
 ```
 
+`--rollout-session-id-scope=group` requires `--router-policy consistent_hashing`. With group scope, conflicting explicit non-empty IDs within one group are rejected. Group affinity may reduce cross-worker load balancing, so it is opt-in. `cache_aware` and `consistent_hashing` are distinct routing policies; group scope only affects the routing key passed to consistent hashing and is not a KV cache manager.
+
 **How it works:**
 
-1. Each sample is assigned a unique `session_id` via UUID
+1. Each sample is assigned a UUID by default, or a rollout group shares one UUID when group scope is selected
 2. On each request, slime passes `X-SMG-Routing-Key: <session_id>` in the HTTP header
 3. SGLang Model Gateway's consistent hashing policy maps this key to a specific worker
-4. Subsequent turns reuse the same `session_id`, ensuring they hit the same worker
+4. Subsequent turns reuse the stored `session_id`, preserving the selected routing identity
 
 ---
 

docs/zh/advanced/sglang-config.md

@@ -301,20 +301,27 @@ slime 会自己启动 router，并把这些外部引擎注册进去。
 
 ### 多轮 Agent 的会话亲和路由
 
-对于多轮对话和 agentic 场景，会话亲和确保同一对话的所有请求路由到同一个 backend worker。这可以显著提升 prefix cache 命中率，因为 worker 已经缓存了对话历史。
+对于多轮对话和 agentic 场景，会话亲和会将共享路由身份的请求路由到同一个 backend worker。这是一种 placement 控制，不保证 cache hit 或性能提升。
 
-slime 自动为每个 sample 分配一个唯一的 `session_id`（存储在 `sample.session_id` 中）。当 router 策略为 `consistent_hashing` 时，该 ID 通过 `X-SMG-Routing-Key` header 传递，SGLang Model Gateway 使用它将同一会话的所有轮次确定性地路由到同一个 worker。
+默认情况下，slime 会为每个 sample 自动分配唯一的 `session_id`（存储在 `sample.session_id` 中）。用户也可以选择 group scope，让同一个 rollout group 内的 sibling samples 共享 routing key，适用于 fan-out 和分组的多轮 rollout。调用方仍可直接设置 `Sample.session_id`；显式设置的非空 ID 会被保留。
 
 ```bash
 --router-policy consistent_hashing
+# 默认：每个 sample 生成一个 session ID。
+--rollout-session-id-scope sample
+
+# 显式启用：每个 rollout group 使用一个生成或继承的 session ID。
+--rollout-session-id-scope group
 ```
 
+`--rollout-session-id-scope=group` 要求 `--router-policy consistent_hashing`。在 group scope 下，同一 group 中冲突的显式非空 ID 会被拒绝。group affinity 可能降低跨 worker 的负载均衡，因此需要显式启用。`cache_aware` 和 `consistent_hashing` 是不同的路由策略；group scope 只影响传给 consistent hashing 的 routing key，并不是 KV cache manager。
+
 **工作原理：**
 
-1. 每个 sample 通过 UUID 分配唯一的 `session_id`
+1. 默认每个 sample 通过 UUID 分配唯一的 `session_id`；选择 group scope 时，一个 rollout group 共享一个 UUID
 2. 每次请求时，slime 在 HTTP header 中传递 `X-SMG-Routing-Key: <session_id>`
 3. SGLang Model Gateway 的 consistent hashing 策略将该 key 映射到特定的 worker
-4. 后续轮次复用相同的 `session_id`，确保命中同一个 worker
+4. 后续轮次复用保存的 `session_id`，保持所选的路由身份
 
 ---
 

slime/rollout/sglang_rollout.py

@@ -40,6 +40,29 @@
 _PROCESSOR_PROMPT_KEYS = {"input_ids", "attention_mask"}
 
 
+def _assign_missing_session_ids(group: list[Sample], scope: str) -> None:
+    """Assign missing routing session IDs according to the requested scope."""
+    if scope == "sample":
+        for sample in group:
+            if sample.session_id is None:
+                sample.session_id = str(uuid.uuid4())
+        return
+
+    if scope != "group":
+        raise ValueError(f"Unsupported rollout session ID scope: {scope!r}.")
+
+    explicit_session_ids = {sample.session_id for sample in group if sample.session_id}
+    if len(explicit_session_ids) > 1:
+        raise ValueError("Group-scoped session IDs require at most one explicit non-empty session_id per group.")
+
+    group_session_id = next(iter(explicit_session_ids), None)
+    if group_session_id is None:
+        group_session_id = str(uuid.uuid4())
+    for sample in group:
+        if not sample.session_id:
+            sample.session_id = group_session_id
+
+
 def _prepare_prompt_ids(sample: Sample, tokenizer, processor: Any) -> list[int]:
     raw_multimodal_inputs = sample.multimodal_inputs or {}
     has_multimodal_inputs = any(value is not None for value in raw_multimodal_inputs.values())
@@ -306,10 +329,7 @@ async def generate_and_rm_group(
     if state.aborted:
         return group
 
-    # Generate a unique session_id for each sample in the group
-    for sample in group:
-        if sample.session_id is None:
-            sample.session_id = str(uuid.uuid4())
+    _assign_missing_session_ids(group, getattr(args, "rollout_session_id_scope", "sample"))
 
     tasks = []
     for idx, sample in enumerate(group):

slime/utils/arguments.py

@@ -425,6 +425,15 @@ def add_rollout_arguments(parser):
                     "This is used to shuffle the prompts and also for the random sampling of the prompts."
                 ),
             )
+            parser.add_argument(
+                "--rollout-session-id-scope",
+                choices=["sample", "group"],
+                default="sample",
+                help=(
+                    "Scope for automatically assigned rollout session IDs. "
+                    "'sample' assigns one ID per sample; 'group' shares one ID within a rollout group."
+                ),
+            )
 
             # sampling
             parser.add_argument(
@@ -1715,6 +1724,10 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]:
 def slime_validate_args(args):
     args.eval_datasets = _resolve_eval_datasets(args)
 
+    if getattr(args, "rollout_session_id_scope", "sample") == "group":
+        if getattr(args, "router_policy", None) != "consistent_hashing":
+            raise ValueError("--rollout-session-id-scope=group requires --router-policy=consistent_hashing.")
+
     if args.kl_coef != 0 or args.use_kl_loss:
         if not os.path.exists(args.ref_load):
             raise FileNotFoundError(f"ref_load {args.ref_load} does not exist, please check the path.")

tests/test_group_session_affinity.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import sys
+from argparse import Namespace
+from contextlib import contextmanager
+from pathlib import Path
+
+import pytest
+
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+sys.path.insert(0, str(Path(__file__).resolve().parent / "plugin_contracts"))
+
+from _shared import install_paths, install_stubs
+
+install_paths()
+install_stubs(with_sglang_router=True, with_transformers=True)
+
+from slime.rollout import sglang_rollout
+from slime.utils.types import Sample
+
+NUM_GPUS = 0
+
+
+def _samples(*session_ids: str | None) -> list[Sample]:
+    return [Sample(index=index, session_id=session_id) for index, session_id in enumerate(session_ids)]
+
+
+@pytest.mark.unit
+def test_sample_scope_assigns_distinct_ids_to_missing_samples():
+    samples = _samples(None, None, None)
+
+    sglang_rollout._assign_missing_session_ids(samples, "sample")
+
+    assert len({sample.session_id for sample in samples}) == 3
+    assert all(sample.session_id for sample in samples)
+
+
+@pytest.mark.unit
+def test_sample_scope_preserves_explicit_session_id():
+    samples = _samples("caller-id", None)
+
+    sglang_rollout._assign_missing_session_ids(samples, "sample")
+
+    assert samples[0].session_id == "caller-id"
+    assert samples[1].session_id and samples[1].session_id != "caller-id"
+
+
+@pytest.mark.unit
+def test_group_scope_assigns_one_id_to_all_missing_samples():
+    samples = _samples(None, None, None)
+
+    sglang_rollout._assign_missing_session_ids(samples, "group")
+
+    assert len({sample.session_id for sample in samples}) == 1
+    assert samples[0].session_id
+
+
+@pytest.mark.unit
+def test_group_scope_inherits_one_explicit_session_id():
+    samples = _samples(None, "caller-id", None)
+
+    sglang_rollout._assign_missing_session_ids(samples, "group")
+
+    assert [sample.session_id for sample in samples] == ["caller-id", "caller-id", "caller-id"]
+
+
+@pytest.mark.unit
+def test_group_scope_rejects_conflicting_explicit_session_ids():
+    with pytest.raises(ValueError, match="at most one explicit non-empty session_id"):
+        sglang_rollout._assign_missing_session_ids(_samples("first", "second"), "group")
+
+
+@pytest.mark.unit
+def test_session_id_scope_rejects_unknown_value():
+    with pytest.raises(ValueError, match="Unsupported rollout session ID scope"):
+        sglang_rollout._assign_missing_session_ids(_samples(None), "rollout")
+
+
+@pytest.mark.unit
+def test_session_id_survives_serialization_and_retry():
+    samples = _samples(None, None)
+
+    sglang_rollout._assign_missing_session_ids(samples, "group")
+    restored = [Sample.from_dict(sample.to_dict()) for sample in samples]
+    before_retry = [sample.session_id for sample in restored]
+    sglang_rollout._assign_missing_session_ids(restored, "group")
+
+    assert [sample.session_id for sample in restored] == before_retry
+
+
+@pytest.mark.unit
+def test_group_scope_preserves_sample_order():
+    samples = _samples(None, None, None)
+    order = [id(sample) for sample in samples]
+
+    sglang_rollout._assign_missing_session_ids(samples, "group")
+
+    assert [id(sample) for sample in samples] == order
+
+
+@pytest.mark.unit
+def test_group_scope_flows_through_group_generate_to_http_with_one_routing_key(monkeypatch):
+    captured = []
+
+    async def fake_post(url, payload, headers=None):
+        captured.append(headers)
+        return {"meta_info": {"finish_reason": {"type": "stop"}}, "text": "response"}
+
+    monkeypatch.setattr(sglang_rollout, "GenerateState", _FakeGenerateState)
+    monkeypatch.setattr(sglang_rollout, "post", fake_post)
+
+    async def fake_rm(args, sample):
+        return 0.0
+
+    monkeypatch.setattr(sglang_rollout, "async_rm", fake_rm)
+    args = Namespace(
+        ci_test=False,
+        rollout_session_id_scope="group",
+        sglang_enable_deterministic_inference=False,
+        group_rm=False,
+        partial_rollout=False,
+        mask_offpolicy_in_partial_rollout=False,
+        custom_generate_function_path=None,
+        sglang_router_ip="127.0.0.1",
+        sglang_router_port=30000,
+        router_policy="consistent_hashing",
+        use_rollout_routing_replay=False,
+        sglang_speculative_algorithm=False,
+    )
+    group = _samples(None, None, None)
+    for sample in group:
+        sample.tokens = [11, 12]
+    sampling_params = {"max_new_tokens": 4, "temperature": 0.0}
+
+    result = asyncio.run(sglang_rollout.generate_and_rm_group(args, group, sampling_params))
+
+    assert result == group
+    assert len(captured) == 3
+    assert {headers["X-SMG-Routing-Key"] for headers in captured} == {group[0].session_id}
+    assert group[0].session_id and all(sample.session_id == group[0].session_id for sample in group)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("scope", ["sample", "group"])
+def test_empty_group_is_a_noop(scope):
+    samples: list[Sample] = []
+
+    sglang_rollout._assign_missing_session_ids(samples, scope)
+
+    assert samples == []
+
+
+@pytest.mark.unit
+def test_default_scope_preserves_per_sample_sampling_seeds(monkeypatch):
+    captured_sampling_params = []
+
+    class FakeGroupGenerateState:
+        def __init__(self, args) -> None:
+            self.aborted = False
+            self.group_sampling_seeds = [101, 102]
+
+    async def fake_generate_and_rm(args, sample, sampling_params, evaluation=False):
+        captured_sampling_params.append(sampling_params)
+        return sample
+
+    monkeypatch.setattr(sglang_rollout, "GenerateState", FakeGroupGenerateState)
+    monkeypatch.setattr(sglang_rollout, "generate_and_rm", fake_generate_and_rm)
+    args = Namespace(
+        rollout_session_id_scope="sample",
+        sglang_enable_deterministic_inference=True,
+        group_rm=False,
+    )
+    group = _samples(None, None)
+    sampling_params = {"temperature": 0.5}
+
+    result = asyncio.run(sglang_rollout.generate_and_rm_group(args, group, sampling_params))
+
+    assert result == group
+    assert len({sample.session_id for sample in group}) == 2
+    assert captured_sampling_params == [
+        {"temperature": 0.5, "sampling_seed": 101},
+        {"temperature": 0.5, "sampling_seed": 102},
+    ]
+    assert sampling_params == {"temperature": 0.5}
+
+
+class _FakeGenerateState:
+    def __init__(self, args) -> None:
+        self.aborted = False
+        self.semaphore = asyncio.Semaphore(3)
+        self.tokenizer = None
+        self.processor = None
+
+    @contextmanager
+    def dp_rank_context(self):
+        yield 0
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "router_policy, expected_headers",
+    [
+        ("consistent_hashing", {"X-SMG-Routing-Key": "caller-id"}),
+        ("round_robin", None),
+        ("cache_aware", None),
+    ],
+)
+def test_generate_sends_routing_header_only_for_consistent_hashing(monkeypatch, router_policy, expected_headers):
+    captured = {}
+
+    async def fake_post(url, payload, headers=None):
+        captured.update(url=url, payload=payload, headers=headers)
+        return {"meta_info": {"finish_reason": {"type": "stop"}}, "text": "response"}
+
+    monkeypatch.setattr(sglang_rollout, "GenerateState", _FakeGenerateState)
+    monkeypatch.setattr(sglang_rollout, "post", fake_post)
+    args = Namespace(
+        ci_test=False,
+        sglang_router_ip="127.0.0.1",
+        sglang_router_port=30000,
+        router_policy=router_policy,
+        use_rollout_routing_replay=False,
+        sglang_speculative_algorithm=False,
+    )
+    sample = Sample(prompt="prompt", tokens=[11, 12], session_id="caller-id")
+    sampling_params = {"max_new_tokens": 4, "temperature": 0.5}
+
+    asyncio.run(sglang_rollout.generate(args, sample, sampling_params))
+
+    assert captured["url"] == "http://127.0.0.1:30000/generate"
+    assert captured["headers"] == expected_headers
+    assert captured["payload"] == {"sampling_params": sampling_params, "return_logprob": True, "input_ids": [11, 12]}
+    assert sample.tokens == [11, 12]
+    assert sampling_params == {"max_new_tokens": 4, "temperature": 0.5}