[Feature] Add consistent hashing routing policy for rollout (#1588)

Author: yitianlian

docs/en/advanced/slime-router.md

@@ -91,3 +91,48 @@ For more details on SGLang Model Gateway, see the [official documentation](https
 - Use SlimeRouter when you need R3 or radix-tree caching
 - Use SGLang Model Gateway for everything else (recommended default)
 
+---
+
+## 4. Session-Affinity Routing for Multi-Turn Agents
+
+When using SGLang Model Gateway with consistent hashing routing policy, Slime automatically assigns each rollout session a unique session ID and uses it as the routing key to enable session affinity.
+
+### What is session affinity?
+
+Session affinity (also called sticky sessions) ensures that all requests belonging to the same conversation or agent session are routed to the same backend worker. This is beneficial for:
+
+- **Multi-turn dialogues**: Keeping the same worker improves prefix cache hit rates
+- **Multi-agent systems**: Ensures agent state consistency and better resource locality
+- **Debugging**: Makes it easier to trace and debug specific sessions
+
+### How it works
+
+When the rollout system generates samples, each sample is assigned a unique `session_id`. This ID is:
+
+1. Automatically generated using UUID for each sample
+2. Stored in `sample.session_id` field
+3. Passed as `X-SMG-Routing-Key` header when the router policy is `consistent_hashing`
+
+The SGLang Model Gateway's consistent hashing policy then uses this routing key to deterministically select the same worker for all requests with the same session ID.
+
+### Configuration
+
+To enable session-affinity routing, simply configure the router policy in Slime:
+
+```bash
+--sglang-router-policy consistent_hashing
+```
+
+Slime will automatically start SGLang Model Gateway with the consistent hashing policy.
+
+> **Note**: If you encounter an error about the `consistent_hashing` policy not being available, upgrade sglang-router:
+> ```bash
+> pip install -U sglang-router
+> ```
+
+### Notes
+
+- Each sample gets its own unique session ID
+- Different samples in the same group may be routed to different workers
+- The same sample's subsequent turns will maintain the same session ID
+- Currently, this feature is only available for SGLang Model Gateway

docs/zh/advanced/slime-router.md

@@ -91,3 +91,48 @@ SGLang Model Gateway 是一个高性能 Rust router，面向大规模 inference
 - 当你需要 R3 或 radix-tree cache 时，使用 SlimeRouter
 - 其他情况使用 SGLang Model Gateway（推荐默认选项）
 
+---
+
+## 4. 多轮 Agent 的会话亲和性路由
+
+当使用 SGLang Model Gateway 的一致性哈希路由策略时，slime 会自动为每个 rollout session 分配唯一的 session ID，并将其作为路由键来实现会话亲和性。
+
+### 什么是会话亲和性？
+
+会话亲和性（Session Affinity，也称为粘性会话）确保属于同一对话或 agent session 的所有请求都被路由到相同的后端 worker。这对以下场景有益：
+
+- **多轮对话**：保持相同 worker 可以提高前缀缓存命中率
+- **多 agent 系统**：确保 agent 状态一致性和更好的资源局部性
+- **调试**：更容易追踪和调试特定会话
+
+### 工作原理
+
+当 rollout 系统生成样本时，每个样本都会被分配一个唯一的 `session_id`。这个 ID 会：
+
+1. 使用 UUID 为每个样本自动生成
+2. 存储在 `sample.session_id` 字段中
+3. 当 router policy 为 `consistent_hashing` 时，作为 `X-SMG-Routing-Key` header 传递
+
+SGLang Model Gateway 的一致性哈希策略会使用这个路由键，确定性地为所有具有相同 session ID 的请求选择相同的 worker。
+
+### 配置方法
+
+启用会话亲和性路由只需在 slime 中配置 router policy 参数：
+
+```bash
+--sglang-router-policy consistent_hashing
+```
+
+Slime 会自动启动 SGLang Model Gateway 并使用一致性哈希策略。
+
+> **注意**：如果遇到 `consistent_hashing` policy 不可用的错误，请升级 sglang-router：
+> ```bash
+> pip install -U sglang-router
+> ```
+
+### 注意事项
+
+- 每个样本都有自己独立的 session ID
+- 同一 group 中的不同样本可能会被路由到不同的 worker
+- 同一样本的后续轮次会保持相同的 session ID
+- 目前该功能只适用于 SGLang Model Gateway

slime/backends/sglang_utils/arguments.py

@@ -19,6 +19,12 @@ def add_sglang_router_arguments(parser):
         default=None,
         help="Port of the SGLang router",
     )
+    parser.add_argument(
+        "--sglang-router-policy",
+        type=str,
+        default=None,
+        help="Routing policy for the SGLang router (e.g., 'consistent_hashing', 'round_robin')",
+    )
     parser.add_argument(
         "--sglang-router-request-timeout-secs",
         type=int,

slime/ray/rollout.py

@@ -658,6 +658,9 @@ def _start_router(args):
         router_args.log_level = "warn"
         router_args.request_timeout_secs = args.sglang_router_request_timeout_secs
 
+        if hasattr(args, "sglang_router_policy") and args.sglang_router_policy:
+            router_args.policy = args.sglang_router_policy
+
         if args.prefill_num_servers is not None:
             router_args.pd_disaggregation = True
 

slime/rollout/sglang_rollout.py

@@ -2,6 +2,7 @@
 import copy
 import inspect
 import logging
+import uuid
 from argparse import Namespace
 from collections.abc import Callable
 from contextlib import contextmanager
@@ -157,7 +158,12 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
         if not sample.tokens:  # Initialize sample.tokens for the first turn
             sample.tokens = prompt_ids
 
-    output = await post(url, payload)
+    # Use session_id for consistent hashing routing if router uses consistent_hashing policy
+    headers = None
+    if args.sglang_router_policy == "consistent_hashing" and sample.session_id:
+        headers = {"X-SMG-Routing-Key": sample.session_id}
+
+    output = await post(url, payload, headers=headers)
 
     if args.use_slime_router and "RadixTreeMiddleware" in args.slime_router_middleware_paths:
         from slime.router.middleware_hub.radix_tree_middleware import postprocess_sample_with_radix_tree
@@ -272,6 +278,11 @@ 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())
+
     tasks = []
     for idx, sample in enumerate(group):
         current_sampling_params = sampling_params.copy()

slime/utils/http_utils.py

@@ -162,12 +162,12 @@ def _next_actor():
     return actor
 
 
-async def _post(client, url, payload, max_retries=60):
+async def _post(client, url, payload, max_retries=60, headers=None):
     retry_count = 0
     while retry_count < max_retries:
         response = None
         try:
-            response = await client.post(url, json=payload or {})
+            response = await client.post(url, json=payload or {}, headers=headers)
             response.raise_for_status()
             content = await response.aread()
             try:
@@ -270,7 +270,7 @@ async def do_post(self, url, payload, max_retries=60):
     _post_actors = created
 
 
-async def post(url, payload, max_retries=60):
+async def post(url, payload, max_retries=60, headers=None):
     # If distributed mode is enabled and actors exist, dispatch via Ray.
     if _distributed_post_enabled and _post_actors:
         try:
@@ -285,7 +285,7 @@ async def post(url, payload, max_retries=60):
             logger.info(f"[http_utils] Distributed POST failed, falling back to local: {e} (url={url})")
             # fall through to local
 
-    return await _post(_http_client, url, payload, max_retries)
+    return await _post(_http_client, url, payload, max_retries, headers=headers)
 
 
 async def get(url):

slime/utils/types.py

@@ -45,6 +45,9 @@ class Status(Enum):
     # metadata used during training, e.g., what loss to use for this sample.
     train_metadata: dict | None = None
 
+    # Session ID for consistent hashing routing (used when router policy is consistent_hashing)
+    session_id: str | None = None
+
     non_generation_time: float = 0.0  # time spent in non-generation steps
 
     @dataclass