diff --git a/docs/en/advanced/slime-router.md b/docs/en/advanced/slime-router.md index 6ec8e17182..592e9ca836 100644 --- a/docs/en/advanced/slime-router.md +++ b/docs/en/advanced/slime-router.md @@ -8,11 +8,10 @@ slime includes an optional slime router used during rollout / data generation. I slime router is a small FastAPI service that: -- Registers workers (SGLang HTTP servers) into a local pool -- Routes requests to a selected worker (simple least-inflight load balancing) -- Proxies arbitrary paths to the selected worker (e.g. `/generate`) +- Registers workers (SGLang HTTP servers) into a local pool, with support for **prefill / decode / regular** worker types +- Routes requests to a selected worker via least-inflight load balancing or **PD dual-dispatch routing** +- Streams proxied responses (e.g. `/generate`) without buffering the full body, improving throughput under high concurrency - Runs periodic health checks and quarantines unhealthy workers -- Supports middleware plugins (via `--slime-router-middleware-paths`) to implement rollout-specific processing (e.g. caching, request/response transforms) In slime's architecture, the router is part of the rollout system ("SGLang + router") that generates samples and pushes them into the data buffer. @@ -27,26 +26,9 @@ In distributed training, slime will start a router automatically when `--sglang- ## 2. Why we need slime router -Unlike production inference, RL rollout needs to capture additional metadata for training: token-level logprobs, loss masks, and (for MoE models) expert routing decisions. slime router provides these capabilities through its middleware system and passthrough proxy design. +Unlike production inference, RL rollout needs to capture additional metadata for training: token-level logprobs, loss masks, and (for MoE models) expert routing decisions. slime router provides these capabilities through its passthrough proxy design. -### 2.1 Radix-tree cache (transparent token management) - -> Use this when your rollout pipeline is text-in/text-out and you cannot reliably persist token IDs; if you already control token-in/token-out (e.g. search r1, multiturn VLM examples), you likely don't need the radix-tree cache. - -Text-in text-out interfaces can cause token retokenization mismatches - re-tokenizing text at training time may produce different token sequences than rollout, breaking per-token alignment needed for PPO/GRPO losses. - -The radix-tree cache solves this transparently: it intercepts text-based requests, tokenizes them, and stores trajectories (text, token IDs, logprobs, loss masks) keyed by the text prefix. After rollout finishes, calling `/retrieve_from_text` returns the exact token sequence with aligned metadata, without requiring any changes to your rollout code. - -Implementation-wise, the radix-tree cache: - -- Accepts text plus tokens/metadata and stores them in a radix tree -- Uses longest-prefix matching to reuse cached token sequences (enabling token-in/token-out downstream) -- Allows insertion of new text continuations as rollout proceeds (multiple trajectories per prompt, e.g. GRPO) -- Periodically cleans up stale nodes to control memory usage - -Use the radix cache when you have text-based rollout code and want token-level precision without rewriting, or when running GRPO with multiple trajectories sharing the same prompt prefix. - -### 2.2 Rollout routing replay (R3) for MoE +### 2.1 Rollout routing replay (R3) for MoE For MoE models, slime supports rollout routing replay (R3): record expert routing decisions during rollout and replay them during training to improve stability. @@ -72,6 +54,17 @@ slime consumes the routing data and replays it during training: We need slime router because the SGLang worker returns routed experts in the response (`meta_info.routed_experts`) when the request sets `return_routed_experts=true`, and slime router preserves this field end-to-end. SGLang Model Gateway may drop this extra metadata when it reconstructs responses with a fixed schema (see section 3). +### 2.2 PD disaggregation + +slime router supports **Prefill-Decode (PD) disaggregation**. When prefill and decode workers are registered, the router automatically enables PD mode: + +- Workers register themselves with a `worker_type` (`prefill`, `decode`, or `regular`) via the `POST /workers` endpoint. +- For each request, the router picks a (prefill, decode) worker pair via least-inflight load balancing, injects bootstrap information (`bootstrap_host`, `bootstrap_port`, `bootstrap_room`) into the request body, and sends the same modified request to **both** workers concurrently. +- The decode worker's response is returned to the caller. The actual KV-cache transfer between workers is coordinated internally via the bootstrap connection. +- If no prefill/decode workers exist, the router falls back to standard single-worker routing. + +This mirrors the dual-dispatch approach used by SGLang Model Gateway's PD router. + --- ## 3. Differences vs SGLang Model Gateway @@ -80,7 +73,7 @@ slime router and SGLang Model Gateway can both route requests to workers, but th ### Key differences -slime router is a lightweight Python/FastAPI proxy that acts as a passthrough to SGLang workers. This passthrough design enables RL-specific features like radix-tree trajectory caching and R3 (which require preserving raw response metadata like `routed_experts`). +slime router is a lightweight Python/FastAPI proxy that acts as a passthrough to SGLang workers. This passthrough design enables RL-specific features like R3 (which require preserving raw response metadata like `routed_experts`). SGLang Model Gateway is a high-performance Rust-based router optimized for large-scale inference: async non-blocking routing, advanced fault tolerance (retries, circuit breakers), multiple load balancing policies (including cache-aware routing), and PD disaggregation support. However, it reconstructs responses with a fixed schema, so it does not preserve the metadata needed for slime's R3 flow. @@ -88,5 +81,5 @@ For more details on SGLang Model Gateway, see the [official documentation](https ### When to use which -- Use slime router when you need R3 or radix-tree caching +- Use slime router when you need R3 or PD disaggregation with metadata preservation - Use SGLang Model Gateway for everything else (recommended default) diff --git a/docs/en/get_started/customization.md b/docs/en/get_started/customization.md index bcbd993ef2..5e352a309f 100644 --- a/docs/en/get_started/customization.md +++ b/docs/en/get_started/customization.md @@ -28,7 +28,6 @@ Below is a summary of all available customization interfaces and their purposes. | [`--custom-megatron-init-path`](#17-megatron-hooks) | Custom initialization after Megatron setup. | | [`--custom-megatron-before-log-prob-hook-path`](#17-megatron-hooks) | Custom logic before log probability computation. | | [`--custom-megatron-before-train-step-hook-path`](#17-megatron-hooks) | Custom logic before each training step. | -| [`--slime-router-middleware-paths`](#18-slime-router-middleware---slime-router-middleware-paths) | Add custom middleware to slime router. | ## Detailed Interface Reference @@ -400,18 +399,7 @@ def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler --- -### 18. slime Router Middleware (`--slime-router-middleware-paths`) - -**Purpose**: Add custom middleware to the slime router for request processing. - -**Use Cases**: -- Request/response transformation -- Custom routing logic -- Caching and optimization - ---- - -### 19. MoE Routing Replay +### 18. MoE Routing Replay Stabilize MoE RL training by recording and replaying expert routing decisions to ensure consistency. diff --git a/docs/zh/advanced/slime-router.md b/docs/zh/advanced/slime-router.md index 20c4cb6e0a..1b6b482b60 100644 --- a/docs/zh/advanced/slime-router.md +++ b/docs/zh/advanced/slime-router.md @@ -8,11 +8,10 @@ slime 提供一个可选的 slime router，用于 rollout / data generation 阶 slime router 是一个小型 FastAPI 服务，主要能力包括： -- 注册 worker（SGLang HTTP server）到本地池 -- 路由请求到选定的 worker（简单的 least-inflight load balancing） -- 代理任意路径到选定的 worker（例如 `/generate`） +- 注册 worker（SGLang HTTP server）到本地池，支持 **prefill / decode / regular** worker 类型 +- 路由请求到选定的 worker——支持 least-inflight 负载均衡和 **PD 双发路由** +- 流式代理请求到选定的 worker（例如 `/generate`），不缓冲完整 response body，提高高并发下的吞吐 - 定期 health checks，并隔离不健康的 worker -- 支持 middleware plugins（通过 `--slime-router-middleware-paths`）实现 rollout 特定处理（例如 caching、request/response transform） 在 slime 架构中，router 是 rollout 系统（"SGLang + router"）的一部分：负责生成样本并将其推入数据缓冲区。 @@ -27,26 +26,9 @@ slime router 是一个小型 FastAPI 服务，主要能力包括： ## 2. 为什么需要 slime router -与 production inference 不同，RL rollout 往往需要捕获用于训练的额外 metadata：token-level logprobs、loss masks，以及（对 MoE 模型）expert routing decisions。slime router 通过 middleware system 和 passthrough proxy 设计提供这些能力。 +与 production inference 不同，RL rollout 往往需要捕获用于训练的额外 metadata：token-level logprobs、loss masks，以及（对 MoE 模型）expert routing decisions。slime router 通过 passthrough proxy 设计提供这些能力。 -### 2.1 Radix-tree cache（透明的 token 管理） - -> 当你的 rollout 流程是 text-in/text-out、并且很难可靠地保存 token IDs 时，适合用 radix-tree cache；如果你已经能自己控制 token-in/token-out（例如 search r1、multiturn VLM 这些 example），通常不需要 radix-tree cache。 - -text-in text-out 接口可能导致 token retokenization mismatches：训练阶段重新 tokenize 文本，得到的 token 序列可能与 rollout 阶段不同，从而破坏 PPO/GRPO 这类方法所需的 per-token alignment。 - -radix-tree cache 可以透明地解决这个问题：它拦截 text-based request，对其进行 tokenize，并将 trajectory（text、token IDs、logprobs、loss masks）按文本前缀作为 key 存储。rollout 结束后，调用 `/retrieve_from_text` 就能取回与 rollout 完全一致的 token 序列及其对齐的 metadata，无需修改现有 rollout 代码。 - -实现上，radix-tree cache 会做几件事： - -- 接收 text 以及 tokens/metadata，并写入 radix tree -- 通过 longest-prefix matching 复用已缓存的 token 序列（使后续流程可以走 token-in/token-out） -- rollout 过程中持续插入新的 text continuation（同一 prompt 下可有多条 trajectory，例如 GRPO） -- 定期清理 stale nodes，控制内存占用 - -当你有 text-based rollout 代码、想获得 token-level 精度但又不想重写，或者在 GRPO 场景中多个 trajectory 共享相同 prompt 前缀时，建议使用 radix-tree cache。 - -### 2.2 Rollout routing replay (R3) for MoE +### 2.1 Rollout routing replay (R3) for MoE 对 MoE 模型，slime 支持 rollout routing replay (R3)：在 rollout 期间记录 expert routing decisions，并在训练期间 replay，以提升训练稳定性。 @@ -72,6 +54,17 @@ slime 侧消费路由数据，并在训练中完成 replay： 我们需要 slime router，是因为当请求设置 `return_routed_experts=true` 时，SGLang worker 会在响应里返回路由信息（`meta_info.routed_experts`），而 slime router 会端到端保留这个字段。SGLang Model Gateway 会用固定 schema 重建响应，可能会丢掉这类额外 metadata（细节见第 3 节）。 +### 2.2 PD 分离（Prefill-Decode 分离） + +slime router 支持 **Prefill-Decode (PD) 分离**。当 prefill 和 decode worker 注册后，router 会自动启用 PD 模式： + +- Worker 注册时携带 `worker_type`（`prefill`、`decode` 或 `regular`），通过 `POST /workers` 端点。 +- 对每个请求，router 通过 least-inflight 负载均衡选择一对 (prefill, decode) worker，向请求体注入 bootstrap 信息（`bootstrap_host`、`bootstrap_port`、`bootstrap_room`），然后将同一个修改后的请求**并发发送**给两个 worker。 +- Decode worker 的响应返回给调用方。实际的 KV-cache 传输由 worker 通过 bootstrap 连接内部协调完成。 +- 如果没有 prefill/decode worker 存在，router 回退到标准的单 worker 路由。 + +这与 SGLang Model Gateway 的 PD router 所使用的双发方式一致。 + --- ## 3. 与 SGLang Model Gateway 的区别 @@ -80,7 +73,7 @@ slime router 与 SGLang Model Gateway 都能将请求路由到 worker，但它 ### 主要区别 -slime router 是一个轻量级的 Python/FastAPI proxy，作为 SGLang worker 的 passthrough proxy。这种 passthrough 设计使得 RL 特定功能成为可能，例如 radix-tree trajectory caching 和 R3（需要保留原始 response metadata，如 `routed_experts`）。 +slime router 是一个轻量级的 Python/FastAPI proxy，作为 SGLang worker 的 passthrough proxy。这种 passthrough 设计使得 RL 特定功能成为可能，例如 R3（需要保留原始 response metadata，如 `routed_experts`）。 SGLang Model Gateway 是一个高性能 Rust router，面向大规模 inference 优化：async non-blocking routing、高级 fault tolerance（retries、circuit breakers）、多种 load balancing policy（包括 cache-aware routing），以及 PD disaggregation 支持。但它会用固定 schema 重建响应，因此不保留 slime 的 R3 流程所需 metadata。 @@ -88,5 +81,5 @@ SGLang Model Gateway 是一个高性能 Rust router，面向大规模 inference ### 如何选择 -- 当你需要 R3 或 radix-tree cache 时，使用 SlimeRouter +- 当你需要 R3 或需要保留 metadata 的 PD 分离时，使用 SlimeRouter - 其他情况使用 SGLang Model Gateway（推荐默认选项） diff --git a/docs/zh/get_started/customization.md b/docs/zh/get_started/customization.md index 5d1e3ddffd..364266e3c3 100644 --- a/docs/zh/get_started/customization.md +++ b/docs/zh/get_started/customization.md @@ -28,7 +28,6 @@ slime 通过函数路径参数提供了广泛的自定义能力。这些参数 | [`--custom-megatron-init-path`](#17-megatron-hook) | Megatron 设置后的自定义初始化。 | | [`--custom-megatron-before-log-prob-hook-path`](#17-megatron-hook) | log probability 计算前的自定义逻辑。 | | [`--custom-megatron-before-train-step-hook-path`](#17-megatron-hook) | 每个训练步骤前的自定义逻辑。 | -| [`--slime-router-middleware-paths`](#18-slime-router-中间件---slime-router-middleware-paths) | 向 slime router 添加自定义中间件。 | ## 详细接口参考 @@ -402,18 +401,7 @@ def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler --- -### 18. slime Router 中间件 (`--slime-router-middleware-paths`) - -**用途**: 向 slime router 添加自定义中间件用于请求处理。 - -**使用场景**: -- 请求/响应转换 -- 自定义路由逻辑 -- 缓存和优化 - ---- - -### 19. MoE 路由重放 +### 18. MoE 路由重放 通过记录和重放专家路由决策来稳定 MoE RL 训练。 diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py index b03a7a5982..0c3430a547 100644 --- a/slime/backends/sglang_utils/sglang_engine.py +++ b/slime/backends/sglang_utils/sglang_engine.py @@ -195,10 +195,8 @@ def _init_normal(self, server_args_dict): return if self.node_rank == 0 and self.router_ip and self.router_port: - if parse(sglang_router.__version__) <= parse("0.2.1") or self.args.use_slime_router: - assert ( - self.worker_type == "regular" - ), "pd disaggregation is not supported in old router or slime router." + if not self.args.use_slime_router and parse(sglang_router.__version__) <= parse("0.2.1"): + assert self.worker_type == "regular", "pd disaggregation is not supported in old router." response = requests.post( f"http://{self.router_ip}:{self.router_port}/add_worker?url=http://{self.server_host}:{self.server_port}" ) @@ -317,7 +315,7 @@ def shutdown(self): if self.worker_type != "encoder" and self.node_rank == 0: worker_url = f"http://{self.server_host}:{self.server_port}" response = None - if parse(sglang_router.__version__) <= parse("0.2.1") or self.args.use_slime_router: + if self.args.use_slime_router or parse(sglang_router.__version__) <= parse("0.2.1"): response = requests.post( f"http://{self.router_ip}:{self.router_port}/remove_worker?url=http://{self.server_host}:{self.server_port}" ) diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 86ecb9064c..33b988eee6 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -912,9 +912,6 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool router_args = copy.copy(args) router_args.sglang_router_ip = router_ip router_args.sglang_router_port = router_port - if has_pd_disaggregation: - router_args.slime_router_pd_disaggregation = True - else: from sglang_router.launch_router import RouterArgs diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py index 392ddfef42..89d9bcff1c 100644 --- a/slime/rollout/sglang_rollout.py +++ b/slime/rollout/sglang_rollout.py @@ -177,37 +177,33 @@ 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 - # Use session_id for consistent hashing routing if router uses consistent_hashing policy + # Use session_id for consistent hashing routing (SGLang Model Gateway) headers = None - if getattr(args, "router_policy", None) == "consistent_hashing" and sample.session_id: - headers = {"X-SMG-Routing-Key": sample.session_id} + if sample.session_id: + if getattr(args, "router_policy", None) == "consistent_hashing": + 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 - - sample = await postprocess_sample_with_radix_tree(args, sample, output) + if "output_token_logprobs" in output["meta_info"]: + new_response_tokens = [item[1] for item in output["meta_info"]["output_token_logprobs"]] + new_response_log_probs = [item[0] for item in output["meta_info"]["output_token_logprobs"]] else: - if "output_token_logprobs" in output["meta_info"]: - new_response_tokens = [item[1] for item in output["meta_info"]["output_token_logprobs"]] - new_response_log_probs = [item[0] for item in output["meta_info"]["output_token_logprobs"]] - else: - new_response_tokens, new_response_log_probs = [], [] + new_response_tokens, new_response_log_probs = [], [] - # Update sample with tokens directly - avoiding re-tokenization - sample.tokens = sample.tokens + new_response_tokens - sample.response_length += len(new_response_tokens) - sample.response += output["text"] + # Update sample with tokens directly - avoiding re-tokenization + sample.tokens = sample.tokens + new_response_tokens + sample.response_length += len(new_response_tokens) + sample.response += output["text"] - # When partial rollout and masking off policy is enabled, update the loss mask - if sample.loss_mask is not None: - assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout - sample.loss_mask += [1] * len(new_response_tokens) + # When partial rollout and masking off policy is enabled, update the loss mask + if sample.loss_mask is not None: + assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout + sample.loss_mask += [1] * len(new_response_tokens) - if sample.rollout_log_probs is None: - sample.rollout_log_probs = [] - sample.rollout_log_probs += new_response_log_probs + if sample.rollout_log_probs is None: + sample.rollout_log_probs = [] + sample.rollout_log_probs += new_response_log_probs if "routed_experts" in output["meta_info"]: sample.rollout_routed_experts = np.frombuffer( diff --git a/slime/router/middleware_hub/__init__.py b/slime/router/middleware_hub/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/slime/router/middleware_hub/radix_tree.py b/slime/router/middleware_hub/radix_tree.py deleted file mode 100644 index 6e722f1e25..0000000000 --- a/slime/router/middleware_hub/radix_tree.py +++ /dev/null @@ -1,681 +0,0 @@ -from __future__ import annotations - -""" -String-based Radix Trie for efficient prefix matching and token caching. -Optimized for string prefixes with corresponding token IDs. -""" - -import threading -import time -from dataclasses import dataclass -from typing import Any - - -@dataclass -class MatchResult: - """Result of prefix matching operation.""" - - matched_prefix: str - token_ids: list[int] - logp: list[float] - loss_mask: list[int] # Added loss mask for model generation parts - remaining_string: str - last_node: StringTreeNode - - -class StringTreeNode: - """Tree node for string-based radix trie.""" - - counter = 0 - - def __init__(self, node_id: int | None = None): - # Core tree structure - self.children: list[StringTreeNode] = [] # Use list to store children - self.parent: StringTreeNode | None = None - - # Node data - self.string_key: str = "" # The string fragment this node represents - self.token_ids: list[int] | None = None # Token IDs for this node only (not cumulative) - self.logp: list[float] | None = None # Log probabilities for this node's tokens - self.loss_mask: list[int] | None = None # Loss mask for model generation parts - - # Access tracking - self.last_access_time = time.monotonic() - self.access_count = 0 - - # Reference counting for protection from eviction - self.ref_count = 0 - - # Weight version tracking - self.weight_version: int | None = None # Weight version for this node - - # Node identification - self.id = StringTreeNode.counter if node_id is None else node_id - StringTreeNode.counter += 1 - - @property - def is_leaf(self) -> bool: - """Check if this node is a leaf node.""" - return len(self.children) == 0 - - @property - def has_value(self) -> bool: - """Check if this node has token IDs stored.""" - return self.token_ids is not None - - def validate_token_logp_consistency(self) -> bool: - """Validate that token_ids, logp, and loss_mask have consistent lengths.""" - if self.token_ids is None and self.logp is None and self.loss_mask is None: - return True - - # Check if at least one is not None - if self.token_ids is not None and len(self.token_ids) > 0: - token_len = len(self.token_ids) - if self.logp is not None and len(self.logp) != token_len: - return False - if self.loss_mask is not None and len(self.loss_mask) != token_len: - return False - - return True - - @property - def is_evictable(self) -> bool: - """Check if this node can be evicted.""" - return self.ref_count == 0 and self.token_ids is not None - - def touch(self): - """Update access time and count.""" - self.last_access_time = time.monotonic() - self.access_count += 1 - - def __lt__(self, other: StringTreeNode) -> bool: - """For heap operations - least recently used first.""" - return self.last_access_time < other.last_access_time - - -class StringRadixTrie: - """ - String-based Radix Trie for efficient prefix matching and token caching. - Features: - - Efficient string prefix matching - - Token ID caching for matched prefixes - - Thread-safe operations - - Weight version tracking - - Automatic garbage collection based on weight version thresholds - """ - - def __init__(self, max_cache_size: int = 10000, gc_threshold_k: int = 5, tokenizer=None, verbose: bool = False): - """ - Initialize the String Radix Trie. - Args: - max_cache_size: Maximum number of cached token IDs (triggers GC when exceeded) - gc_threshold_k: GC threshold - nodes with weight_version < (current_version - k) will be removed - tokenizer: Optional tokenizer for converting text to tokens when not found in cache - verbose: Whether to print debug information and tree structure - """ - self.max_cache_size = max_cache_size - self.gc_threshold_k = gc_threshold_k - self.tokenizer = tokenizer - self.verbose = verbose - - # Tree structure - self.root = StringTreeNode() - self.root.string_key = "" - self.root.ref_count = 1 # Root is always protected - - # Cache statistics - self.total_entries = 0 - self.cache_hits = 0 - self.cache_misses = 0 - self.cur_cache_size = 0 # Total number of token IDs across all nodes - - # Thread safety - self._lock = threading.RLock() - - def find_longest_prefix(self, text: str) -> MatchResult: - """ - Find the longest cached prefix for the given text. - Args: - text: Input string to find prefix for - Returns: - MatchResult containing matched prefix, token IDs, logp, and remaining string - """ - with self._lock: - if not text: - return MatchResult("", [], [], [], text, self.root) - - matched_tokens = [] - matched_logp = [] - matched_loss_mask = [] - matched_prefix = "" - current_node = self.root - remaining_text = text - - while remaining_text: - # Find the best matching child that completely matches from start - best_child = None - best_key_len = 0 - - for child_node in current_node.children: - # Only consider complete startswith matches using node's string_key - if remaining_text.startswith(child_node.string_key): - if len(child_node.string_key) > best_key_len: - best_child = child_node - best_key_len = len(child_node.string_key) - - if best_child is None: - # No complete startswith match found - break - - # Move to the best matching child - best_child.touch() - current_node = best_child - matched_prefix += best_child.string_key - remaining_text = remaining_text[best_key_len:] - - # Accumulate tokens, logp, and loss_mask from this node - if best_child.has_value: - matched_tokens.extend(best_child.token_ids) - matched_logp.extend(best_child.logp) - if best_child.loss_mask is not None: - matched_loss_mask.extend(best_child.loss_mask) - else: - # If no loss_mask is stored, create default mask same as logp - matched_loss_mask.extend([1] * len(best_child.token_ids)) - self.cache_hits += 1 - - if not matched_tokens: - self.cache_misses += 1 - - result = MatchResult( - matched_prefix, matched_tokens, matched_logp, matched_loss_mask, remaining_text, current_node - ) - - # Print tree structure if verbose is enabled - if self.verbose: - print("Tree structure after find_longest_prefix:") - self.pretty_print() - - return result - - def insert( - self, - text: str, - token_ids: list[int], - logp: list[float] | None = None, - loss_mask: list[int] | None = None, - weight_version: int | None = None, - ) -> bool: - """ - Insert a string and its corresponding token IDs, log probabilities, and loss mask into the trie. - Args: - text: String to insert - token_ids: Corresponding token IDs - logp: Corresponding log probabilities (must match token_ids length) - loss_mask: Corresponding loss mask for model generation parts (must match token_ids length) - weight_version: Optional weight version for this insertion - Returns: - True if insertion was successful - """ - with self._lock: - if not text or not token_ids: - if self.verbose: - print("[RadixTree] Insertion failed: text or token_ids is empty") - return False - - # Use provided weight version - current_weight_version = weight_version - - # Validate logp consistency - if logp is not None and len(logp) != len(token_ids): - if self.verbose: - print( - f"[WARNING] Logp length {len(logp)} does not match token length {len(token_ids)} for text: {text}" - ) - print(f"[WARNING] Logp: {logp}") - print(f"[WARNING] Token IDs: {token_ids}") - return False - - # Validate loss_mask consistency - if loss_mask is not None and len(loss_mask) != len(token_ids): - if self.verbose: - print( - f"[WARNING] Loss mask length {len(loss_mask)} does not match token length {len(token_ids)} for text: {text}" - ) - print(f"[WARNING] Loss mask: {loss_mask}") - print(f"[WARNING] Token IDs: {token_ids}") - return False - - # If logp is not provided, create default values (0.0) - if logp is None: - logp = [0.0] * len(token_ids) - - # If loss_mask is not provided, create default values (1 for model generation parts) - if loss_mask is None: - loss_mask = [0] * len(token_ids) - - result = self._insert(text, token_ids, logp, loss_mask, current_weight_version) - - # Check if GC should be triggered after insert - if self.cur_cache_size > self.max_cache_size and weight_version is not None: - if self.verbose: - print( - f"[RadixTree] Cache size {self.cur_cache_size} exceeds limit {self.max_cache_size}, triggering GC" - ) - gc_removed = self.gc_by_weight_version(weight_version) - if self.verbose: - print(f"[RadixTree] GC removed {gc_removed} nodes, new cache size: {self.cur_cache_size}") - - # Print tree structure if verbose is enabled - if self.verbose: - print("Tree structure after insert:") - self.pretty_print() - - return result - - def _insert( - self, - text: str, - token_ids: list[int], - logp: list[float], - loss_mask: list[int], - weight_version: int | None = None, - ) -> bool: - """Insert tokens - skip tokens for existing nodes just like we skip text.""" - - current_node = self.root - remaining_text = text - remaining_tokens = token_ids[:] # Copy the tokens list - remaining_logp = logp[:] # Copy the logp list - remaining_loss_mask = loss_mask[:] # Copy the loss_mask list - - # Track all nodes traversed during insert for weight version update - traversed_nodes = [current_node] - new_node = None - - while remaining_text: - # Find best startswith match - best_child = None - best_key_len = 0 - - for child_node in current_node.children: - if remaining_text.startswith(child_node.string_key) and len(child_node.string_key) > best_key_len: - best_child = child_node - best_key_len = len(child_node.string_key) - - if best_child is not None: - # Found existing node - skip its text and tokens - current_node = best_child - traversed_nodes.append(current_node) - remaining_text = remaining_text[best_key_len:] - - # Skip the tokens that this existing node covers - if best_child.has_value: - tokens_to_skip = len(best_child.token_ids) - remaining_tokens = remaining_tokens[tokens_to_skip:] - remaining_logp = remaining_logp[tokens_to_skip:] - remaining_loss_mask = remaining_loss_mask[tokens_to_skip:] - else: - # Create new node for remaining text with remaining tokens - new_node = StringTreeNode() - new_node.parent = current_node - new_node.string_key = remaining_text - - if remaining_tokens: # Only assign if there are tokens left - new_node.token_ids = remaining_tokens - new_node.logp = remaining_logp - new_node.loss_mask = remaining_loss_mask - new_node.touch() - # Increment cache size by number of tokens added - self.cur_cache_size += len(remaining_tokens) - - current_node.children.append(new_node) - traversed_nodes.append(new_node) - self.total_entries += 1 - break - - # If we've traversed the entire text and the last node doesn't have tokens, - # assign remaining tokens to it - if remaining_text == "" and not current_node.has_value: - if remaining_tokens: # Only assign if there are tokens left - current_node.token_ids = remaining_tokens - current_node.logp = remaining_logp - current_node.loss_mask = remaining_loss_mask - current_node.touch() - self.cur_cache_size += len(remaining_tokens) - - # Update weight version for all traversed nodes - if weight_version is not None and new_node: - new_node.weight_version = weight_version - - return True - - def remove(self, text: str) -> bool: - """ - Remove a string and all nodes with this text as prefix from the trie. - Args: - text: String to remove (will also remove all strings starting with this text) - Returns: - True if any removal was performed - """ - with self._lock: - node = self._find_node_by_text(text) - if node: - removed_count = self._clean_node_subtree(node) - - # Print tree structure if verbose is enabled - if self.verbose: - print("Tree structure after remove:") - self.pretty_print() - - return removed_count > 0 - return False - - def _find_node_by_text(self, text: str) -> StringTreeNode | None: - """ - Find node by exact text match. - Args: - text: Text to find - Returns: - Node if found, None otherwise - """ - result = self.find_longest_prefix(text) - if result.matched_prefix == text: - return result.last_node - return None - - def _clean_node_subtree(self, node: StringTreeNode) -> int: - """ - Clean a node and all its descendants. - This is the core cleanup function. - Args: - node: Node to clean (including all descendants) - Returns: - Number of nodes removed - """ - if node == self.root: - return 0 - return self._remove_node_and_descendants(node) - - def _remove_node_and_descendants(self, node: StringTreeNode) -> int: - """ - Remove a node and all its descendants from the trie. - Args: - node: The node to remove along with all its descendants - Returns: - Number of nodes removed - """ - if node == self.root: - # Never remove root node - return 0 - - removed_count = 0 - - # First, recursively remove all descendants - for child in list(node.children): # Create a copy to avoid modification during iteration - removed_count += self._remove_node_and_descendants(child) - - # Count this node if it has data and decrement cache size - if node.has_value: - removed_count += 1 - # Decrement cache size by number of tokens removed - self.cur_cache_size -= len(node.token_ids) - - # Remove this node from its parent - if self._remove_node_from_parent(node): - # Update count for the node structure itself - pass # _remove_node_from_parent already decrements total_entries - - return removed_count - - def _remove_node_from_parent(self, node: StringTreeNode) -> bool: - """Remove a node from its parent's children list.""" - if node.parent and node in node.parent.children: - node.parent.children.remove(node) - self.total_entries -= 1 - return True - return False - - def gc_by_weight_version(self, current_weight_version: int | None = None) -> int: - """ - Perform garbage collection based on weight version. - Remove nodes with weight_version < (current_weight_version - gc_threshold_k). - Args: - current_weight_version: Current weight version to use for GC threshold - Returns: - Number of nodes removed - """ - with self._lock: - if current_weight_version is None: - if self.verbose: - print("[RadixTree GC] No weight version provided, skipping GC") - return 0 - - gc_threshold = current_weight_version - self.gc_threshold_k - if self.verbose: - print( - f"[RadixTree GC] Starting GC with threshold: {gc_threshold} (current_version: {current_weight_version}, k: {self.gc_threshold_k})" - ) - - nodes_to_remove = self._find_outdated_nodes(gc_threshold) - removed_count = 0 - - for node in nodes_to_remove: - # Validate that subtree weight versions are <= parent weight version - self._validate_subtree_weight_versions(node) - removed_count += self._clean_node_subtree(node) - - if self.verbose: - print(f"[RadixTree GC] Completed GC, removed {removed_count} nodes") - - return removed_count - - def _find_outdated_nodes(self, gc_threshold: int) -> list[StringTreeNode]: - """ - Find nodes that should be removed based on weight version threshold. - Uses layer-by-layer traversal - if parent is outdated, children are not checked. - Args: - gc_threshold: Weight version threshold (nodes < this value will be removed) - Returns: - List of nodes to remove - """ - outdated_nodes = [] - - def check_node(node): - if node == self.root: - # Root is never removed, check its children - for child in node.children: - check_node(child) - return - - # Check if this node should be removed - if node.weight_version is not None and node.weight_version <= gc_threshold and node.has_value: - outdated_nodes.append(node) - return # Don't check children since entire subtree will be removed - - # Node is not outdated, check its children - for child in node.children: - check_node(child) - - check_node(self.root) - return outdated_nodes - - def _validate_subtree_weight_versions(self, node: StringTreeNode): - """ - Validate that all nodes in subtree have weight_version <= parent weight_version. - Args: - node: Root node of subtree to validate - """ - - def validate_recursive(current_node, parent_weight_version): - if current_node.weight_version is not None and parent_weight_version is not None: - assert current_node.weight_version <= parent_weight_version, ( - f"Child node weight_version {current_node.weight_version} > " - f"parent weight_version {parent_weight_version}" - ) - - # Recursively validate children - for child in current_node.children: - validate_recursive(child, current_node.weight_version) - - # Start validation from the node itself - validate_recursive(node, node.weight_version) - - def get_stats(self) -> dict[str, Any]: - """Get cache statistics.""" - with self._lock: - total_requests = self.cache_hits + self.cache_misses - hit_rate = self.cache_hits / total_requests if total_requests > 0 else 0 - - return { - "total_entries": self.total_entries, - "cache_hits": self.cache_hits, - "cache_misses": self.cache_misses, - "hit_rate": hit_rate, - "max_cache_size": self.max_cache_size, - "cur_cache_size": self.cur_cache_size, - "gc_threshold_k": self.gc_threshold_k, - } - - def clear(self): - """Clear all entries from the trie.""" - with self._lock: - self.root = StringTreeNode() - self.root.string_key = "" - self.root.ref_count = 1 - self.total_entries = 0 - self.cache_hits = 0 - self.cache_misses = 0 - self.cur_cache_size = 0 - - def pretty_print(self): - """Print the trie structure in a readable format.""" - print("String Radix Trie Structure:") - print("=" * 50) - self._print_node(self.root, 0) - print("=" * 50) - stats = self.get_stats() - for key, value in stats.items(): - print(f"{key}: {value}") - - def _print_node(self, node: StringTreeNode, depth: int): - """Recursively print node structure.""" - indent = " " * depth - key_repr = repr(node.string_key) if node.string_key else "" - token_info = "" - if node.has_value: - token_info = f" -> tokens: {node.token_ids}" - if node.logp: - token_info += f", logp: {[round(p, 3) for p in node.logp]}" - if node.loss_mask: - token_info += f", loss_mask: {node.loss_mask}" - access_info = f" (accessed: {node.access_count}, ref: {node.ref_count})" - - print(f"{indent}{key_repr}{token_info}{access_info}") - - for child in node.children: - self._print_node(child, depth + 1) - - def retrieve_from_text(self, text: str, return_logprob: bool = True): - """ - Get tokens from text by looking up in radix tree or using tokenizer. - Also fetches weight version from worker during this operation. - Args: - text: Input text to get tokens for - return_logprob: If True, also return log probabilities - Returns: - List of token IDs corresponding to the input text if return_logprob is False. - Tuple of (token_ids, logp) if return_logprob is True. - """ - # Call find_longest_prefix to get the match result - result = self.find_longest_prefix(text) - - # If we have a match and it covers the entire text, return the tokens - if result.matched_prefix and result.token_ids: - additional_tokens = self.tokenizer(result.remaining_string, add_special_tokens=False)["input_ids"] - return ( - result.token_ids + additional_tokens, - ( - result.logp + len(additional_tokens) * [0.0] - if return_logprob - else [0] * len(result.token_ids + additional_tokens) - ), - result.loss_mask + len(additional_tokens) * [0], - ) - # If result is empty and input text is not empty, tokenize with tokenizer - # This is needed because we cannot get the prompt token id from engine response - # We have to manually insert the text and token into the tree - if self.tokenizer and text: - # Tokenize the text using the provided tokenizer - tokens = self.tokenizer(text, add_special_tokens=False)["input_ids"] - # Insert the text and tokens into the tree - self.insert(text, tokens) - # Return the tokens - return (tokens, [0.0] * len(tokens), [0] * len(tokens)) - else: - raise ValueError("Tokenizer or input text can't be empty") - - -# Example usage and testing -if __name__ == "__main__": - # Create trie instance for testing - trie = StringRadixTrie(max_cache_size=100, verbose=True) - - # Example usage with simplified insert - test_cases = [ - ("Hello world", [1, 2, 3], [-0.1, -0.2, -0.3]), - ("Hello", [1, 2], [-0.1, -0.2]), - ("Hi there", [4, 5, 6], [-0.4, -0.5, -0.6]), - ] - - # Insert test data with weight version and loss masks - print("Inserting test data...") - for text, tokens, logp in test_cases: - # Create loss_mask to match tokens length, 1 for model generation parts - loss_mask = [1] * len(tokens) - success = trie.insert(text, tokens, logp, loss_mask, weight_version=1) - print(f"Inserted '{text}' -> {tokens}: {success}") - - print("\nTrie structure:") - trie.pretty_print() - - # Test prefix matching - print("\nTesting prefix matching:") - test_queries = [ - "Hello world!", # Should match "Hello world" completely - "Hello everyone", # Should match "Hello" only - "Hi there", # Should match "Hi" only - "How are you doing?", # Should match "How are you" completely - "Goodbye", # Should not match anything - "Hell", # Should not match anything (not complete startswith) - ] - - for query in test_queries: - result = trie.find_longest_prefix(query) - print(f"Query: '{query}'") - print( - f" Matched: '{result.matched_prefix}' -> tokens: {result.token_ids}, logp: {result.logp}, loss_mask: {result.loss_mask}" - ) - print(f" Remaining: '{result.remaining_string}'") - print() - - # Test removal - print("Testing removal:") - removed = trie.remove("Hello") - print(f"Removed 'Hello': {removed}") - - result = trie.find_longest_prefix("Hello world") - print( - f"After removal - 'Hello world' -> matched: '{result.matched_prefix}', tokens: {result.token_ids}, logp: {result.logp}, loss_mask: {result.loss_mask}" - ) - - # Show final stats - print("\nFinal statistics:") - stats = trie.get_stats() - for key, value in stats.items(): - print(f"{key}: {value}") - - # Test GC with weight version - print("\nTesting GC with weight version 5:") - gc_removed = trie.gc_by_weight_version(5) - print(f"GC removed {gc_removed} nodes") diff --git a/slime/router/middleware_hub/radix_tree_middleware.py b/slime/router/middleware_hub/radix_tree_middleware.py deleted file mode 100644 index 5df163a38d..0000000000 --- a/slime/router/middleware_hub/radix_tree_middleware.py +++ /dev/null @@ -1,169 +0,0 @@ -import asyncio -import json - -from fastapi.responses import JSONResponse -from starlette.middleware.base import BaseHTTPMiddleware -from starlette.requests import Request -from starlette.responses import Response -from transformers import AutoTokenizer - -from slime.utils.http_utils import post -from slime.utils.mask_utils import get_response_lengths -from slime.utils.types import Sample - -from .radix_tree import StringRadixTrie - -# Hop-by-hop headers that should not be forwarded -HOP_BY_HOP = { - "content-length", - "transfer-encoding", - "connection", - "keep-alive", - "proxy-authenticate", - "proxy-authorization", - "te", - "trailers", - "upgrade", -} - - -def _filter_headers(headers): - """Filter out hop-by-hop headers that should not be forwarded.""" - return {k: v for k, v in headers.items() if k.lower() not in HOP_BY_HOP} - - -async def _materialize_response(resp): - """Convert streaming-like Response into a regular Response/JSONResponse safely.""" - # Collect all bytes from the streaming response - body = b"" - async for chunk in resp.body_iterator: - body += chunk - - # Try to parse as JSON based on content-type - ct = resp.headers.get("content-type", "") - headers = _filter_headers(resp.headers) - - if "application/json" in ct: - # If it's JSON, try to parse and return as JSONResponse - try: - data = json.loads(body.decode("utf-8")) - return JSONResponse(content=data, status_code=resp.status_code, headers=headers) - except Exception: - # JSON parsing failed, fall back to raw bytes - pass - - # Other types: return as raw bytes (without content-length) - return Response(content=body, status_code=resp.status_code, headers=headers, media_type=resp.media_type) - - -class RadixTreeMiddleware(BaseHTTPMiddleware): - def __init__(self, app, *, router): - super().__init__(app) - self.router = router - self.args = router.args - self.tokenizer = AutoTokenizer.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True) - self.radix_tree = StringRadixTrie(max_cache_size=10000, tokenizer=self.tokenizer, verbose=False) - self.router.radix_tree = self.radix_tree - - async def dispatch(self, request: Request, call_next): - - path = request.url.path - - if path != "/generate": - return await call_next(request) - - request_json = await request.json() - if "text" in request_json: - input_text = request_json.pop("text", "") - elif "input_ids" in request_json: - input_text = self.tokenizer.decode(request_json["input_ids"]) - else: - input_text = None - if not input_text: - return await call_next(request) - input_tokens, input_logprobs, input_loss_mask = self.radix_tree.retrieve_from_text( - input_text, return_logprob=True - ) - request_json["input_tokens"] = input_tokens - request_json["stream"] = False - request._json = request_json - - response_data = None - for _ in range(5): - response = await call_next(request) - - # If upstream returned a streaming response, materialize it to avoid Content-Length issues - if response.__class__.__name__ == "_StreamingResponse": - response = await _materialize_response(response) - # Try to parse JSON from the current response for meta inspection - try: - if hasattr(response, "body") and isinstance(response.body, (bytes, bytearray)): - response_data = json.loads(response.body.decode("utf-8")) - elif hasattr(response, "content") and isinstance(response.content, (dict, list)): - response_data = response.content # JSONResponse.content is already a dict/list - except Exception: - response_data = None - - if ( - isinstance(response_data, dict) - and "meta_info" in response_data - and "finish_reason" in response_data["meta_info"] - and response_data["meta_info"]["finish_reason"]["type"] != "abort" - ): - break - # await 30 seconds for aborted responses - await asyncio.sleep(30) - - if isinstance(response_data, dict) and "text" in response_data and "output_ids" in response_data: - generated_text = response_data["text"] - - full_text = input_text + generated_text - if full_text: - try: - if "output_token_logprobs" in response_data.get("meta_info", {}): - generated_token_logprobs = [ - item[0] for item in response_data["meta_info"]["output_token_logprobs"] - ] - generated_token_ids = [item[1] for item in response_data["meta_info"]["output_token_logprobs"]] - full_logprobs = input_logprobs + generated_token_logprobs - full_token_ids = input_tokens + generated_token_ids - full_loss_mask = input_loss_mask + [1] * len(generated_token_ids) - self.radix_tree.insert( - full_text, - full_token_ids, - full_logprobs, - full_loss_mask, - weight_version=response_data["meta_info"]["weight_version"], - ) - else: - generated_token_ids = self.tokenizer(generated_text, add_special_tokens=False)["input_ids"] - full_token_ids = input_tokens + generated_token_ids - full_loss_mask = input_loss_mask + [1] * len(generated_token_ids) - self.radix_tree.insert( - full_text, - full_token_ids, - None, - full_loss_mask, - weight_version=response_data["meta_info"]["weight_version"], - ) - - if getattr(self.router, "verbose", False): - print(f"[slime-router] Successfully cached trajectory with {len(full_token_ids)} tokens") - except Exception as e: - if getattr(self.router, "verbose", False): - print(f"[slime-router] Warning: Failed to cache trajectory: {e}") - return response - - -async def postprocess_sample_with_radix_tree(args, sample: Sample, output: dict): - assert not args.partial_rollout, "Currently partial rollout is not supported when using slime router" - retrieve_url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/retrieve_from_text" - retrieve_payload = {"text": sample.prompt + output["text"], "return_logp": True} - retrieve_output = await post(retrieve_url, retrieve_payload) - sample.tokens = retrieve_output["tokens"] - sample.response += output["text"] - sample.loss_mask = retrieve_output["loss_mask"] - sample.response_length = get_response_lengths([sample.loss_mask])[0] - sample.loss_mask = sample.loss_mask[-sample.response_length :] - sample.rollout_log_probs = retrieve_output["rollout_logp"][-sample.response_length :] - return sample diff --git a/slime/router/router.py b/slime/router/router.py index 669094e442..bde130b037 100644 --- a/slime/router/router.py +++ b/slime/router/router.py @@ -2,81 +2,106 @@ import asyncio import json import logging +import uuid +from enum import Enum +from urllib.parse import urlparse import httpx import uvicorn from fastapi import FastAPI, Request -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, StreamingResponse from starlette.responses import Response -from slime.utils.misc import load_function logger = logging.getLogger(__name__) +class WorkerType(str, Enum): + REGULAR = "regular" + PREFILL = "prefill" + DECODE = "decode" + + +class WorkerInfo: + """Metadata for a registered worker.""" + + __slots__ = ("url", "worker_type", "active_requests", "consecutive_failures", "bootstrap_port") + + def __init__(self, url: str, worker_type: WorkerType = WorkerType.REGULAR, bootstrap_port: int | None = None): + self.url = url + self.worker_type = worker_type + self.active_requests: int = 0 + self.consecutive_failures: int = 0 + self.bootstrap_port = bootstrap_port + + def run_router(args): - """ - Run the Slime router with the specified configuration. - """ - # Initialize the router with tokenizer and lazy worker initialization + """Run the Slime router with the specified configuration.""" slime_router = SlimeRouter(args, verbose=False) - - # Start the server uvicorn.run(slime_router.app, host=args.sglang_router_ip, port=args.sglang_router_port, log_level="info") class SlimeRouter: def __init__(self, args, verbose=False): - """Initialize the slime-router with SGLang router address""" + """Initialize the slime-router.""" self.args = args self.verbose = verbose self.app = FastAPI() self.app.add_event_handler("startup", self._start_background_health_check) - # URL -> Active Request Count (load state) - self.worker_request_counts: dict[str, int] = {} - # URL -> Consecutive Failures - self.worker_failure_counts: dict[str, int] = {} + # URL -> WorkerInfo + self.workers: dict[str, WorkerInfo] = {} # Quarantined workers excluded from routing pool self.dead_workers: set[str] = set() self.max_weight_version = None + # --- Connection pool --- max_connections = getattr(args, "slime_router_max_connections", None) if max_connections is None: max_connections = ( args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine ) + # Generous keep-alive pool for high concurrency + max_keepalive = max(max_connections // 2, 20) timeout = getattr(args, "slime_router_timeout", None) self.client = httpx.AsyncClient( - limits=httpx.Limits(max_connections=max_connections), + limits=httpx.Limits( + max_connections=max_connections, + max_keepalive_connections=max_keepalive, + keepalive_expiry=30, + ), timeout=httpx.Timeout(timeout), + http2=True, ) self._setup_routes() - for middleware_path in args.slime_router_middleware_paths or []: - if self.verbose: - print(f"[slime-router] Loading middleware from: {middleware_path}") - middleware = load_function(middleware_path) - self.app.add_middleware(middleware, router=self) + # ------------------------------------------------------------------ + # Routes + # ------------------------------------------------------------------ def _setup_routes(self): - """Setup all the HTTP routes""" - # sglang-router api + """Setup all the HTTP routes.""" self.app.post("/add_worker")(self.add_worker) + self.app.post("/remove_worker")(self.remove_worker) + self.app.post("/workers")(self.add_worker_v2) + self.app.get("/workers")(self.list_workers_v2) self.app.get("/list_workers")(self.list_workers) - self.app.post("/retrieve_from_text")(self.retrieve_from_text) - # Catch-all route for proxying to SGLang - must be registered LAST + self.app.get("/health")(self.health) + # Catch-all route for proxying — must be registered LAST self.app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])(self.proxy) + # ------------------------------------------------------------------ + # Health check background loop + # ------------------------------------------------------------------ + async def _start_background_health_check(self): asyncio.create_task(self._health_check_loop()) - async def _check_worker_health(self, url): - """Encapsulated health check logic for better maintainability.""" + async def _check_worker_health(self, url: str): try: response = await self.client.get(f"{url}/health", timeout=5.0) if response.status_code == 200: @@ -95,31 +120,27 @@ async def _health_check_loop(self): try: await asyncio.sleep(interval) - urls = [u for u in self.worker_request_counts if u not in self.dead_workers] + urls = [u for u in self.workers if u not in self.dead_workers] if not urls: continue results = await asyncio.gather(*(self._check_worker_health(url) for url in urls)) for url, is_healthy in results: + if url not in self.workers: + continue if not is_healthy: - failures = self.worker_failure_counts.get(url, 0) + 1 - self.worker_failure_counts[url] = failures - - if failures >= threshold: + self.workers[url].consecutive_failures += 1 + if self.workers[url].consecutive_failures >= threshold: logger.warning( f"[slime-router] Worker {url} failed {threshold} consecutive health checks. Marking as DEAD." ) self.dead_workers.add(url) - # TODO (chenyang): Connect back 'dead' workers requires a mechanism to sync - # model versions to avoid off-policy issues from stale weights, since these - # dead workers' parameters may not be refitted. else: - self.worker_failure_counts[url] = 0 + self.workers[url].consecutive_failures = 0 - logger.debug( - f"[slime-router] Health check complete. {len(self.worker_request_counts) - len(self.dead_workers)} workers healthy." - ) + alive = sum(1 for u in self.workers if u not in self.dead_workers) + logger.debug(f"[slime-router] Health check complete. {alive} workers healthy.") except asyncio.CancelledError: logger.warning("[slime-router] Background health check loop is being cancelled.") @@ -128,122 +149,251 @@ async def _health_check_loop(self): logger.error(f"[slime-router] Unexpected error in health check loop: {e}", exc_info=True) await asyncio.sleep(5) + # ------------------------------------------------------------------ + # Worker selection + # ------------------------------------------------------------------ + + def _healthy_workers(self, worker_type: WorkerType | None = None) -> list[WorkerInfo]: + """Return live workers, optionally filtered by type.""" + workers = [w for url, w in self.workers.items() if url not in self.dead_workers] + if worker_type is not None: + workers = [w for w in workers if w.worker_type == worker_type] + return workers + + def _select_by_least_inflight(self, candidates: list[WorkerInfo]) -> WorkerInfo: + """Pick the worker with the fewest active requests.""" + if not candidates: + raise RuntimeError("No healthy workers available in the pool") + return min(candidates, key=lambda w: w.active_requests) + + def _is_pd_mode(self) -> bool: + """Check if PD disaggregation is active (prefill workers exist).""" + return any( + w.worker_type == WorkerType.PREFILL for url, w in self.workers.items() if url not in self.dead_workers + ) + + def _pick_pd_pair(self) -> tuple[WorkerInfo, WorkerInfo]: + """Pick a (prefill, decode) worker pair using least-inflight.""" + prefill_candidates = self._healthy_workers(WorkerType.PREFILL) + decode_candidates = self._healthy_workers(WorkerType.DECODE) + if not prefill_candidates: + raise RuntimeError("No healthy prefill workers available") + if not decode_candidates: + raise RuntimeError("No healthy decode workers available") + prefill = self._select_by_least_inflight(prefill_candidates) + decode = self._select_by_least_inflight(decode_candidates) + prefill.active_requests += 1 + decode.active_requests += 1 + return prefill, decode + + def _pick_worker(self) -> WorkerInfo: + """Pick a single worker via least-inflight (non-PD mode).""" + candidates = self._healthy_workers() + worker = self._select_by_least_inflight(candidates) + worker.active_requests += 1 + return worker + + def _finish_worker(self, worker: WorkerInfo): + """Mark the request to the given worker as finished.""" + worker.active_requests -= 1 + assert worker.active_requests >= 0, f"Worker {worker.url} active_requests went negative" + + # ------------------------------------------------------------------ + # Proxy (streaming) + # ------------------------------------------------------------------ + async def proxy(self, request: Request, path: str): - """Proxy all other requests to the SGLang router""" - # Forward all other paths to SGLang router - worker_url = self._use_url() - url = f"{worker_url}/{path}" + """Stream-proxy requests to a selected backend worker. - # Get request body and headers + In PD disaggregation mode, picks a (prefill, decode) pair, injects + bootstrap info, and sends the same request to both workers concurrently + (mirroring sgl-model-gateway behaviour). The decode worker's response + is returned to the caller. + """ body = await request.body() headers = dict(request.headers) + if self._is_pd_mode(): + return await self._proxy_pd(path, body, headers) + else: + worker = self._pick_worker() + try: + return await self._forward_to_worker(worker, path, body, headers) + finally: + self._finish_worker(worker) + + # --- PD dual-dispatch helpers --- + + def _bootstrap_host_from_url(self, worker_url: str) -> str: + """Extract the hostname from a worker URL for bootstrap.""" + return urlparse(worker_url).hostname or "127.0.0.1" + + def _inject_bootstrap(self, body: bytes, prefill: WorkerInfo) -> bytes: + """Inject bootstrap_host / bootstrap_port / bootstrap_room into the request body.""" try: - response = await self.client.request(request.method, url, content=body, headers=headers) - # Eagerly read content so we can return JSON (not streaming) + payload = json.loads(body) if body else {} + except Exception: + return body + payload["bootstrap_host"] = self._bootstrap_host_from_url(prefill.url) + payload["bootstrap_port"] = prefill.bootstrap_port + payload["bootstrap_room"] = uuid.uuid4().hex + return json.dumps(payload).encode() + + async def _proxy_pd(self, path: str, body: bytes, headers: dict) -> Response: + """PD dual dispatch: send the same request to prefill + decode concurrently.""" + prefill, decode = self._pick_pd_pair() + try: + modified_body = self._inject_bootstrap(body, prefill) + + prefill_url = f"{prefill.url}/{path}" + decode_url = f"{decode.url}/{path}" + + prefill_req = self.client.build_request("POST", prefill_url, content=modified_body, headers=headers) + decode_req = self.client.build_request("POST", decode_url, content=modified_body, headers=headers) + + # Fire both concurrently; we only care about the decode response. + _prefill_task = asyncio.ensure_future(self.client.send(prefill_req, stream=True)) + decode_response = await self.client.send(decode_req, stream=True) + + return await self._build_response(decode_response) + finally: + self._finish_worker(prefill) + self._finish_worker(decode) + + async def _forward_to_worker(self, worker: WorkerInfo, path: str, body: bytes, headers: dict) -> Response: + """Forward a request to a single worker and return its response.""" + url = f"{worker.url}/{path}" + req = self.client.build_request("POST", url, content=body, headers=headers) + response = await self.client.send(req, stream=True) + return await self._build_response(response) + + async def _build_response(self, response: httpx.Response) -> Response: + """Convert an httpx streaming response into a FastAPI response.""" + content_type = response.headers.get("content-type", "") + + if "text/event-stream" not in content_type: content = await response.aread() - content_type = response.headers.get("content-type", "") + await response.aclose() try: - # Prefer parsing JSON if possible data = json.loads(content) - return JSONResponse( - content=data, - status_code=response.status_code, - headers=dict(response.headers), - ) + return JSONResponse(content=data, status_code=response.status_code) except Exception: - # Fall back to raw body with original content type - return Response( - content=content, - status_code=response.status_code, - headers=dict(response.headers), - media_type=content_type or None, - ) + return Response(content=content, status_code=response.status_code, media_type=content_type or None) + + async def _stream(): + try: + async for chunk in response.aiter_bytes(): + yield chunk finally: - if response is not None: - await response.aclose() + await response.aclose() - finally: - self._finish_url(worker_url) + return StreamingResponse(_stream(), status_code=response.status_code, media_type=content_type) + + # ------------------------------------------------------------------ + # Worker management endpoints + # ------------------------------------------------------------------ async def add_worker(self, request: Request): - """Add a new worker to the router. - Supports providing the URL via query string or JSON body. + """Add a new worker (v1 compat — query string or JSON body). + Examples: - - POST /add_worker?url=http://127.0.0.1:10090 - - POST /add_worker with body {"url": "http://127.0.0.1:10090"} + POST /add_worker?url=http://127.0.0.1:10090 + POST /add_worker?url=http://127.0.0.1:10090&worker_type=prefill + POST /add_worker {"url": "...", "worker_type": "prefill"} """ - # 1) Prefer query param worker_url = request.query_params.get("url") or request.query_params.get("worker_url") + worker_type_str = request.query_params.get("worker_type", "regular") - # 2) Fallback to JSON body if not worker_url: body = await request.body() payload = json.loads(body) if body else {} worker_url = payload.get("url") or payload.get("worker_url") + worker_type_str = payload.get("worker_type", worker_type_str) if not worker_url: return JSONResponse( - status_code=400, content={"error": "worker_url is required (use query ?url=... or JSON body)"} + status_code=400, content={"error": "url is required (use query ?url=... or JSON body)"} ) - # Add if new, keep a simple request count per worker - if worker_url not in self.worker_request_counts: - self.worker_request_counts[worker_url] = 0 - self.worker_failure_counts[worker_url] = 0 + try: + worker_type = WorkerType(worker_type_str) + except ValueError: + worker_type = WorkerType.REGULAR + + if worker_url not in self.workers: + self.workers[worker_url] = WorkerInfo(url=worker_url, worker_type=worker_type) if self.verbose: - print(f"[slime-router] Added new worker: {worker_url}") + print(f"[slime-router] Added new worker: {worker_url} (type={worker_type.value})") - return {"status": "success", "worker_urls": self.worker_request_counts} + return {"status": "success", "worker_urls": {u: w.active_requests for u, w in self.workers.items()}} - async def list_workers(self, request: Request): - """List all registered workers""" - return {"urls": list(self.worker_request_counts.keys())} + async def add_worker_v2(self, request: Request): + """Add worker — SGLang Model Gateway compatible ``POST /workers`` endpoint. - async def retrieve_from_text(self, request: Request): - """Get token information from text input""" + Body: {"url": "...", "worker_type": "prefill"|"decode"|"regular", "bootstrap_port": 12345} + """ body = await request.body() payload = json.loads(body) if body else {} + worker_url = payload.get("url") + worker_type_str = payload.get("worker_type", "regular") + bootstrap_port = payload.get("bootstrap_port") - text = payload.get("text", "") + if not worker_url: + return JSONResponse(status_code=400, content={"error": "url is required in JSON body"}) - # Use radix tree's retrieve_from_text method (no need to fetch weight version here) - token_ids, logp, loss_mask = self.radix_tree.retrieve_from_text(text, return_logprob=True) + try: + worker_type = WorkerType(worker_type_str) + except ValueError: + worker_type = WorkerType.REGULAR - # Handle the result based on whether logp was requested - result = { - "tokens": token_ids, # token IDs - "response": text, # The input text - "loss_mask": loss_mask, # Loss mask for the tokens - "token_length": len(token_ids), - "loss_mask_length": len(loss_mask), - "rollout_logp": logp, - } + if worker_url not in self.workers: + self.workers[worker_url] = WorkerInfo( + url=worker_url, worker_type=worker_type, bootstrap_port=bootstrap_port + ) + if self.verbose: + print(f"[slime-router] Added new worker: {worker_url} (type={worker_type.value})") - return result + return {"status": "success"} - def _use_url(self): - """Select worker URL with minimal active requests.""" + async def remove_worker(self, request: Request): + """Remove a worker from the pool.""" + worker_url = request.query_params.get("url") or request.query_params.get("worker_url") - if not self.dead_workers: - # Healthy path: select from all workers - url = min(self.worker_request_counts, key=self.worker_request_counts.get) - else: - # Degraded path: select from workers not in dead_workers - valid_workers = (w for w in self.worker_request_counts if w not in self.dead_workers) - try: - url = min(valid_workers, key=self.worker_request_counts.get) - except ValueError: - raise RuntimeError("No healthy workers available in the pool") from None + if not worker_url: + body = await request.body() + payload = json.loads(body) if body else {} + worker_url = payload.get("url") or payload.get("worker_url") + + if not worker_url: + return JSONResponse(status_code=400, content={"error": "url is required"}) - self.worker_request_counts[url] += 1 - return url + self.workers.pop(worker_url, None) + self.dead_workers.discard(worker_url) + return {"status": "success"} - def _finish_url(self, url): - """Mark the request to the given URL as finished""" - assert url in self.worker_request_counts, f"URL {url} not recognized" - self.worker_request_counts[url] -= 1 - assert self.worker_request_counts[url] >= 0, f"URL {url} count went negative" + async def list_workers(self, request: Request): + """List all registered workers (v1 compat).""" + return {"urls": list(self.workers.keys())} + + async def list_workers_v2(self, request: Request): + """List workers — SGLang Model Gateway compatible ``GET /workers``.""" + workers_list = [] + for url, w in self.workers.items(): + entry = { + "url": url, + "worker_type": w.worker_type.value, + "active_requests": w.active_requests, + "is_healthy": url not in self.dead_workers, + } + if w.bootstrap_port is not None: + entry["bootstrap_port"] = w.bootstrap_port + workers_list.append(entry) + return {"workers": workers_list} + + async def health(self, request: Request): + """Router health check endpoint.""" + alive = sum(1 for u in self.workers if u not in self.dead_workers) + return {"status": "ok", "healthy_workers": alive, "total_workers": len(self.workers)} if __name__ == "__main__": @@ -252,10 +402,7 @@ def _finish_url(self, url): parser.add_argument("--port", type=int, default=30000) parser.add_argument("--sglang-host", type=str, required=True) parser.add_argument("--sglang-port", type=int, required=True) - parser.add_argument("--tokenizer-name", type=str, help="Name of the tokenizer to use for tokenization") parser.add_argument("--verbose", action="store_true", help="Enable verbose output") args = parser.parse_args() - - # Run the router run_router(args) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index d8da881ea3..1d4fa24d33 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -1007,12 +1007,6 @@ def add_router_arguments(parser): default=False, help="Whether to use SlimeRouter for text-based routing instead of SGLang token-based routing", ) - parser.add_argument( - "--slime-router-middleware-paths", - type=str, - nargs="+", - default="", - ) parser.add_argument( "--slime-router-timeout", type=float,