diff --git a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py index f4ea1c4f8..ddaf80744 100644 --- a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py +++ b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py @@ -1,5 +1,6 @@ from __future__ import annotations +import dataclasses import logging import os from typing import TYPE_CHECKING @@ -101,6 +102,9 @@ def handle_pd_disaggregation(server_args: ServerArgs) -> None: server_args.disaggregation_transfer_backend != "fake" ), "Prefill server does not support 'fake' as the transfer backend" + if envs.SGLANG_RUST_SERVER.get(): + _alias_bootstrap_port_to_api_port(server_args) + if server_args.disaggregation_mode in ("prefill", "decode"): if ( envs.SGLANG_DISAGG_STAGING_BUFFER.get() @@ -111,3 +115,36 @@ def handle_pd_disaggregation(server_args: ServerArgs) -> None: f"disaggregation_transfer_backend='mooncake' or 'nixl', " f"got '{server_args.disaggregation_transfer_backend}'." ) + + +def _alias_bootstrap_port_to_api_port(server_args: ServerArgs) -> None: + """Rust-server prefill serves the KV bootstrap registry on the api listener + itself, so the resolved bootstrap port must BE the api port — every internal + consumer (KVManager registration, PrefillBootstrapQueue) reads the resolved + field and agrees automatically. Decode is untouched: there the field names + the PREFILL side's bootstrap port and must stay as the operator set it. + """ + default_port = next( + f.default + for f in dataclasses.fields(server_args) + if f.name == "disaggregation_bootstrap_port" + ) + if server_args.disaggregation_bootstrap_port not in ( + default_port, + server_args.port, + ): + raise ValueError( + "SGLANG_RUST_SERVER serves the PD KV bootstrap registry on the api " + "port itself; --disaggregation-bootstrap-port " + f"{server_args.disaggregation_bootstrap_port} conflicts with --port " + f"{server_args.port}. Drop --disaggregation-bootstrap-port (decode " + "nodes and the PD router must then target the prefill api port)." + ) + if server_args.disaggregation_bootstrap_port != server_args.port: + logger.info( + "SGLANG_RUST_SERVER: KV bootstrap registry is served on the api " + "port; disaggregation_bootstrap_port %d -> %d", + server_args.disaggregation_bootstrap_port, + server_args.port, + ) + server_args.disaggregation_bootstrap_port = server_args.port diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 88b51f204..ed4ffb817 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -2251,10 +2251,6 @@ def _execute_server_warmup(server_args: ServerArgs): _global_state.tokenizer_manager.server_status = ServerStatus.Up else: - # TODO: @rainj-me fix this when Rust server supports disaggregation - assert ( - not envs.SGLANG_RUST_SERVER.get() - ), "Rust server is not supported for disaggregation warmup for now" logger.info(f"Start of pd disaggregation warmup ...") status_codes = asyncio.run( _send_disaggregation_warmup_requests( @@ -2272,14 +2268,20 @@ def _execute_server_warmup(server_args: ServerArgs): server_args.dp_size, ) logger.info("End of disaggregation warmup") - _global_state.tokenizer_manager.server_status = ServerStatus.Up else: logger.info( "Disaggregation warmup failed (mode=%s), status codes: %s", server_args.disaggregation_mode, failed_status_codes, ) - _global_state.tokenizer_manager.server_status = ServerStatus.UnHealthy + # In rust-server mode there is no TokenizerManager (readiness is + # the Rust server's own /health), so skip the status update. + if not envs.SGLANG_RUST_SERVER.get(): + _global_state.tokenizer_manager.server_status = ( + ServerStatus.Up + if not failed_status_codes + else ServerStatus.UnHealthy + ) except Exception: last_traceback = get_exception_traceback() diff --git a/python/sglang/srt/managers/disagg_service.py b/python/sglang/srt/managers/disagg_service.py index 8b02add20..406888bc5 100644 --- a/python/sglang/srt/managers/disagg_service.py +++ b/python/sglang/srt/managers/disagg_service.py @@ -27,18 +27,30 @@ def start_disagg_service( host=server_args.host, port=server_args.disaggregation_bootstrap_port, ) - is_create_store = ( - server_args.node_rank == 0 and transfer_backend == TransferBackend.ASCEND + maybe_create_ascend_config_store( + server_args=server_args, transfer_backend=transfer_backend ) - if is_create_store: - try: - from memfabric_hybrid import create_config_store - - ascend_url = os.getenv("ASCEND_MF_STORE_URL") - create_config_store(ascend_url) - except Exception as e: - error_message = f"Failed create mf store, invalid ascend_url." - error_message += f" With exception {e}" - raise error_message return bootstrap_server + + +def maybe_create_ascend_config_store( + server_args: ServerArgs, transfer_backend: TransferBackend +) -> None: + """Also called directly by the rust-server scheduler: there the KV + bootstrap registry is served by the embedded rust server's api listener + (one rust implementation covers every transfer backend — their + bootstrap-server subclasses are all plain ``CommonKVBootstrapServer``, + which the rust registry ports verbatim), leaving this store as the only + ``start_disagg_service`` duty left to perform.""" + if not (server_args.node_rank == 0 and transfer_backend == TransferBackend.ASCEND): + return + try: + from memfabric_hybrid import create_config_store + + ascend_url = os.getenv("ASCEND_MF_STORE_URL") + create_config_store(ascend_url) + except Exception as e: + raise RuntimeError( + f"Failed create mf store, invalid ascend_url. With exception {e}" + ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f2272447d..589d65232 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -99,6 +99,7 @@ from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config from sglang.srt.layers.quantization.unquant import initialize_bf16_gemm_config from sglang.srt.lora.lora_drainer import LoRADrainer from sglang.srt.lora.lora_overlap_loader import LoRAOverlapLoader +from sglang.srt.managers.disagg_service import maybe_create_ascend_config_store from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator from sglang.srt.managers.io_struct import ( AbortReq, @@ -590,6 +591,15 @@ class Scheduler( # Init profiler self.init_profiler() + # Start the embedded Rust frontend (rank 0). Must precede + # init_disaggregation: on PD prefill the rust api listener also serves + # the KV bootstrap registry, and the KVManagers built there register to + # it synchronously. (The listener is bound synchronously inside launch, + # so the registry is accepting once this returns.) Must also precede + # the request receiver, which reads self.recv_from_tokenizer to pick + # its ingress transport. + self.maybe_init_rust_server() + # Init prefill-decodedisaggregation self.init_disaggregation() @@ -618,10 +628,6 @@ class Scheduler( self.maybe_init_scripted_scheduler_hook() - # Start the embedded Rust frontend (rank 0) before the request receiver, - # which reads self.rust_ring_recv to pick its ingress transport. - self.maybe_init_rust_server() - self.init_request_receiver() self.init_dp_attn_adapter() @@ -1214,6 +1220,21 @@ class Scheduler( get_disagg().disaggregation_transfer_backend ) + # In rust-server mode the KV bootstrap registry is already serving on + # the rust api listener (maybe_init_rust_server runs before this + # method — the PrefillBootstrapQueue's KVManager below registers to it + # synchronously, and a failed registration only retries ~60s then logs, + # leaving every PD request unroutable). Only the ascend config store, + # which start_disagg_service would otherwise create, is left to do. + if ( + self.disaggregation_mode == DisaggregationMode.PREFILL + and self._hosts_rust_server() + ): + maybe_create_ascend_config_store( + server_args=self.server_args, + transfer_backend=self.transfer_backend, + ) + # todo: should we fix this when enabling mtp or it doesn't matter since we only enable mtp in decode node thus we don't transfer draft kvs between P and D? draft_token_to_kv_pool = kv_cache_builder.get_draft_kv_pool( draft_worker=self.draft_worker, @@ -1844,17 +1865,22 @@ class Scheduler( else: self.scripted_scheduler_hook = None + def _hosts_rust_server(self) -> bool: + """Whether this scheduler rank embeds the Rust server (rank 0 only) — + and with it the server-process duties a Python ``TokenizerManager`` + would otherwise own (e.g. serving the PD KV bootstrap registry).""" + return envs.SGLANG_RUST_SERVER.get() and ( + self.ps.pp_rank == 0 + and self.ps.attn_tp_rank == 0 + and self.ps.attn_cp_rank == 0 + ) + def maybe_init_rust_server(self) -> None: """Start the embedded Rust server (rank 0) if ``SGLANG_RUST_SERVER`` is set, and point the ingress receiver at it. All the plumbing lives in ``RustServer`` (scheduler_components/rust_scheduler.py).""" - is_rank_zero = ( - self.ps.pp_rank == 0 - and self.ps.attn_tp_rank == 0 - and self.ps.attn_cp_rank == 0 - ) - if not (envs.SGLANG_RUST_SERVER.get() and is_rank_zero): + if not self._hosts_rust_server(): # Always define the attribute: init_output_streamer and the # process_input_requests hook read self.rust_server unconditionally. self.rust_server = None @@ -2303,9 +2329,10 @@ class Scheduler( f"bootstrap room id. {req.rid=}" ) logger.error(error_msg) - recv_req.time_stats.trace_ctx.abort( - abort_info={"reason": error_msg} - ) + if not envs.SGLANG_RUST_SERVER.get(): + recv_req.time_stats.trace_ctx.abort( + abort_info={"reason": error_msg} + ) prepare_abort(req, error_msg, status_code=HTTPStatus.BAD_REQUEST) self.output_streamer.stream_output([req], req.return_logprob) return diff --git a/rust/sglang-server/src/api_server.rs b/rust/sglang-server/src/api_server.rs index 911989b29..2674f3560 100644 --- a/rust/sglang-server/src/api_server.rs +++ b/rust/sglang-server/src/api_server.rs @@ -9,6 +9,7 @@ mod guard; mod log; mod native_api; mod openai; +mod pd_bootstrap; mod submit; use std::sync::Arc; @@ -48,7 +49,7 @@ pub async fn serve( egress_activity, }; // Each endpoint module registers its own routes and merges here. - let app = Router::new() + let mut app = Router::new() .merge(common::routes()) .merge(native_api::routes()) .merge(openai::routes()) @@ -59,6 +60,14 @@ pub async fn serve( // No body limit, matching the Python server. .layer(axum::extract::DefaultBodyLimit::disable()) .with_state(state); + if server_args.enable_pd_bootstrap() { + // Merged after `with_state` (the registry carries its own state) and + // before `log::apply`, so bootstrap traffic shows in the access log. + let (bootstrap_routes, sweeper) = pd_bootstrap::router_and_sweeper(); + tokio::spawn(sweeper); // cancelled with the runtime on shutdown + app = app.merge(bootstrap_routes); + tracing::info!("PD KV bootstrap registry mounted on the api listener"); + } let app = log::apply(app, &server_args); // The listener was already bound synchronously in `runtime::start` (so a port diff --git a/rust/sglang-server/src/api_server/frame.rs b/rust/sglang-server/src/api_server/frame.rs index 70b266818..129f7c619 100644 --- a/rust/sglang-server/src/api_server/frame.rs +++ b/rust/sglang-server/src/api_server/frame.rs @@ -188,11 +188,14 @@ pub(super) fn frame_value(out: &ChunkEvent, rid: &str) -> serde_json::Value { let Some(ex) = out.extras.as_deref() else { return v; }; - if !ex.out_lp_val.is_empty() { + // Python (`add_logprob_to_meta_info`) always sets input+output token + // logprobs together, empty lists included. A PD decode node never receives + // input logprobs (they belong to prefill), yet its response must still + // carry the key — the PD router keys its merge of prefill's + // `input_token_logprobs` on its presence. + if !ex.out_lp_val.is_empty() || !ex.in_lp_val.is_empty() { v["meta_info"]["output_token_logprobs"] = logprob_tuples(&ex.out_lp_val, &ex.out_lp_idx, opt_texts(&ex.out_lp_txt)); - } - if !ex.in_lp_val.is_empty() { v["meta_info"]["input_token_logprobs"] = logprob_tuples(&ex.in_lp_val, &ex.in_lp_idx, opt_texts(&ex.in_lp_txt)); } @@ -269,7 +272,12 @@ pub(super) fn cumulative_frame_json( if let Some(v) = &acc.in_tid_json { let _ = write!(m, ",\"input_token_ids_logprobs\":{v}"); } - if let Some(v) = &acc.in_lp_json { + // Input+output token logprobs are emitted as a PAIR whenever either side has + // data (empty list included), matching `frame_value` byte for byte — see the + // PD-router rationale there. + let lp_pair = acc.in_lp_json.is_some() || !acc.out_lp_json.is_empty(); + if lp_pair { + let v = acc.in_lp_json.as_deref().unwrap_or("[]"); let _ = write!(m, ",\"input_token_logprobs\":{v}"); } if let Some(v) = &acc.in_top_json { @@ -280,7 +288,7 @@ pub(super) fn cumulative_frame_json( if !acc.out_tid_json.is_empty() { let _ = write!(m, ",\"output_token_ids_logprobs\":[{}]", acc.out_tid_json); } - if !acc.out_lp_json.is_empty() { + if lp_pair { let _ = write!(m, ",\"output_token_logprobs\":[{}]", acc.out_lp_json); } if !acc.out_top_json.is_empty() { diff --git a/rust/sglang-server/src/api_server/native_api.rs b/rust/sglang-server/src/api_server/native_api.rs index 4b1d0bee1..4037f46ea 100644 --- a/rust/sglang-server/src/api_server/native_api.rs +++ b/rust/sglang-server/src/api_server/native_api.rs @@ -61,6 +61,10 @@ fn health_routes() -> Router { .route("/health_generate", probe) } +/// Sentinel host that makes the KV connector no-op. Parity with +/// `sglang.srt.disaggregation.utils.FAKE_BOOTSTRAP_HOST`. +const FAKE_BOOTSTRAP_HOST: &str = "2.2.2.2"; + /// `GET /health_generate` — deep health: confirm the scheduler → detok path is /// producing output. 200 iff the egress heartbeat advances within `timeout` /// (from `SGLANG_HEALTH_CHECK_TIMEOUT`, frozen at router build), else 503. @@ -80,6 +84,10 @@ async fn health_generate(State(state): State, timeout: std::time::Dura // Fire the probe (the heartbeat is the signal, not its own response). A busy // scheduler skips it with no terminal frame, so its detok registration is // cleaned up only by the `AbortGuard` below. + // + // On a PD node the scheduler 400-aborts room-less requests, so inject the + // same fake bootstrap pair Python uses (`FAKE_BOOTSTRAP_HOST` / room 0). + let pd = state.server_args.is_disaggregation(); let probe = GenerateRequest { // The `HEALTH_CHECK_` rid form rid: Rid::new_health_check(), @@ -91,6 +99,8 @@ async fn health_generate(State(state): State, timeout: std::time::Dura ..Default::default() }, stream: false, + bootstrap_host: pd.then(|| FAKE_BOOTSTRAP_HOST.into()), + bootstrap_room: pd.then_some(0), ..Default::default() }; let (rid, _keepalive) = diff --git a/rust/sglang-server/src/api_server/pd_bootstrap.rs b/rust/sglang-server/src/api_server/pd_bootstrap.rs new file mode 100644 index 000000000..e369c768d --- /dev/null +++ b/rust/sglang-server/src/api_server/pd_bootstrap.rs @@ -0,0 +1,612 @@ +//! PD KV bootstrap registry — the rust port of Python +//! `CommonKVBootstrapServer` (disaggregation/common/conn.py), which every +//! transfer backend (mooncake / mori / nixl / ascend) subclasses without +//! overrides, so this one implementation covers them all. +//! +//! Prefill ranks PUT their `{rank_ip, rank_port}` to `/route`; decode ranks +//! GET per-rank routes and the aggregate topology (the all `-1` sentinel +//! query); the PD router registers/queries per-room dp ranks. The wire +//! protocol is Python-owned — field names, status codes, and the `-1` +//! sentinel below are parity pins, not this crate's design. +//! +//! Served on the api listener itself: `api_server::serve` merges +//! [`router_and_sweeper`]'s routes on every prefill server +//! (`ServerArgs::enable_pd_bootstrap()`). In rust-server mode the resolved +//! `disaggregation_bootstrap_port` is aliased to the api port, so KV managers +//! register here without knowing about the merge. The scheduler starts the +//! rust server BEFORE `init_disaggregation` — the KV managers register +//! synchronously there, with only a few bounded retries, so the routes must +//! already be accepting. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use axum::extract::{Query, State}; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::routing::{post, put}; +use axum::{Json, Router}; + +/// Python default: `SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL = EnvInt(120)`. +const ENTRY_CLEANUP_INTERVAL_ENV: &str = "SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL"; +const ENTRY_CLEANUP_INTERVAL_DEFAULT_SECS: u64 = 120; + +/// One registered prefill rank, exactly the JSON the Python decode side +/// consumes (`PrefillRankInfo`). +#[derive(Clone, serde::Serialize)] +struct RankInfo { + rank_ip: String, + rank_port: i64, +} + +/// The all-`-1` sentinel response, exactly Python's +/// `dataclasses.asdict(PrefillServerInfo)` shape. +#[derive(serde::Serialize)] +struct PrefillServerInfo { + attn_tp_size: i64, + attn_cp_size: i64, + dp_size: i64, + pp_size: i64, + page_size: Option, + kv_cache_dtype: Option, + follow_bootstrap_room: bool, + enable_dsa_cache_layer_split: bool, + prefill_http_port: Option, +} + +struct RoomEntry { + dp_rank: i64, + registered_at: Instant, +} + +/// Mirror of the Python server's mutable state. First PUT wins for the +/// topology scalars (Python's `if self.x is None` pattern); `registered_count` +/// counts raw PUTs, and readiness is `count >= dp*cp*tp*pp` — both verbatim +/// from Python, re-registrations included. +#[derive(Default)] +struct Registry { + attn_tp_size: Option, + attn_cp_size: Option, + dp_size: Option, + pp_size: Option, + page_size: Option, + kv_cache_dtype: Option, + follow_bootstrap_room: Option, + enable_dsa_cache_layer_split: Option, + prefill_http_port: Option, + /// Keyed `(dp_group, attn_cp_rank, attn_tp_rank, pp_rank)` — the flat form + /// of Python's nested `prefill_port_table` dicts. + prefill_ranks: HashMap<(i64, i64, i64, i64), RankInfo>, + room_to_dp_rank: HashMap, + registered_count: i64, +} + +impl Registry { + /// `dp * cp * tp * pp` once every size is known (saturating: absurd sizes + /// stay "never ready" instead of overflowing). + fn expected(&self) -> Option { + Some( + self.dp_size? + .saturating_mul(self.attn_cp_size?) + .saturating_mul(self.attn_tp_size?) + .saturating_mul(self.pp_size?), + ) + } + + fn is_ready(&self) -> bool { + self.expected() + .is_some_and(|expected| self.registered_count >= expected) + } +} + +type Shared = Arc>; + +/// i64 that also accepts a numeric string — used exactly where Python coerces +/// with `int(data[...])`, so the wire stays as tolerant as the original. +#[derive(Clone, Copy)] +struct Int(i64); + +impl<'de> serde::Deserialize<'de> for Int { + fn deserialize>(d: D) -> Result { + #[derive(serde::Deserialize)] + #[serde(untagged)] + enum Raw { + Int(i64), + Str(String), + } + match Raw::deserialize(d)? { + Raw::Int(v) => Ok(Int(v)), + // `int(...)` tolerates surrounding whitespace. + Raw::Str(s) => s + .trim() + .parse() + .map(Int) + .map_err(|_| serde::de::Error::custom(format!("invalid int: {s:?}"))), + } + } +} + +/// PUT /route payload (`CommonKVManager.register_to_bootstrap`). +#[derive(serde::Deserialize)] +struct RoutePut { + attn_tp_size: i64, + attn_tp_rank: i64, + attn_cp_size: i64, + attn_cp_rank: i64, + attn_dp_size: i64, + attn_dp_rank: i64, + pp_size: i64, + pp_rank: i64, + system_dp_size: i64, + system_dp_rank: i64, + rank_ip: String, + rank_port: Int, + page_size: Int, + #[serde(default)] + kv_cache_dtype: Option, + #[serde(default)] + prefill_http_port: Option, + #[serde(default)] + load_balance_method: Option, + #[serde(default)] + enable_dsa_cache_layer_split: Option, +} + +fn not_ready(registered_count: i64) -> Response { + ( + StatusCode::SERVICE_UNAVAILABLE, + format!("Prefill server not fully registered yet ({registered_count} workers registered)."), + ) + .into_response() +} + +async fn route_put(State(state): State, Json(body): Json) -> Response { + // `system_dp_size == 1` → attention-dp topology; else system-dp topology. + let dp_size = if body.system_dp_size == 1 { + body.attn_dp_size + } else { + body.system_dp_size + }; + let dp_group = if body.system_dp_size == 1 { + body.attn_dp_rank + } else { + body.system_dp_rank + }; + + let mut reg = state.lock().unwrap(); + reg.attn_tp_size.get_or_insert(body.attn_tp_size); + reg.attn_cp_size.get_or_insert(body.attn_cp_size); + reg.dp_size.get_or_insert(dp_size); + reg.pp_size.get_or_insert(body.pp_size); + reg.page_size.get_or_insert(body.page_size.0); + if reg.kv_cache_dtype.is_none() { + reg.kv_cache_dtype = body.kv_cache_dtype; + } + if reg.prefill_http_port.is_none() { + reg.prefill_http_port = body.prefill_http_port.map(|p| p.0); + } + reg.follow_bootstrap_room.get_or_insert( + body.load_balance_method + .as_deref() + .unwrap_or("follow_bootstrap_room") + == "follow_bootstrap_room", + ); + reg.enable_dsa_cache_layer_split + .get_or_insert(body.enable_dsa_cache_layer_split.unwrap_or(false)); + + reg.prefill_ranks.insert( + (dp_group, body.attn_cp_rank, body.attn_tp_rank, body.pp_rank), + RankInfo { + rank_ip: body.rank_ip.clone(), + rank_port: body.rank_port.0, + }, + ); + reg.registered_count += 1; + + tracing::debug!( + dp_group, + cp = body.attn_cp_rank, + tp = body.attn_tp_rank, + pp = body.pp_rank, + rank_ip = %body.rank_ip, + rank_port = body.rank_port.0, + registered = reg.registered_count, + expected = reg.expected(), + "registered prefill bootstrap rank" + ); + "OK".into_response() +} + +async fn route_get( + State(state): State, + Query(query): Query>, +) -> Response { + // A missing, empty (Python truthiness), or non-integer param → 400. + let rank = |k: &str| query.get(k).and_then(|v| v.trim().parse::().ok()); + let (Some(dp), Some(cp), Some(tp), Some(pp)) = ( + rank("prefill_dp_rank"), + rank("prefill_cp_rank"), + rank("target_tp_rank"), + rank("target_pp_rank"), + ) else { + return ( + StatusCode::BAD_REQUEST, + "Missing inputs for bootstrap server.", + ) + .into_response(); + }; + + let reg = state.lock().unwrap(); + // Python checks readiness in both branches; hoisted, same behavior. + if !reg.is_ready() { + return not_ready(reg.registered_count); + } + + if (dp, cp, tp, pp) == (-1, -1, -1, -1) { + // Aggregate-topology sentinel. The sizes are Some — `is_ready` above + // requires all four. + return Json(PrefillServerInfo { + attn_tp_size: reg.attn_tp_size.unwrap(), + attn_cp_size: reg.attn_cp_size.unwrap(), + dp_size: reg.dp_size.unwrap(), + pp_size: reg.pp_size.unwrap(), + page_size: reg.page_size, + kv_cache_dtype: reg.kv_cache_dtype.clone(), + follow_bootstrap_room: reg.follow_bootstrap_room.unwrap_or(true), + enable_dsa_cache_layer_split: reg.enable_dsa_cache_layer_split.unwrap_or(false), + prefill_http_port: reg.prefill_http_port, + }) + .into_response(); + } + + match reg.prefill_ranks.get(&(dp, cp, tp, pp)) { + Some(info) => Json(info.clone()).into_response(), + None => ( + StatusCode::NOT_FOUND, + format!( + "Bootstrap info not found for dp_rank={dp} cp_rank={cp} \ + tp_rank={tp} pp_rank={pp}" + ), + ) + .into_response(), + } +} + +#[derive(serde::Deserialize)] +struct RegisterDpRank { + bootstrap_room: Int, + dp_rank: Int, +} + +async fn register_dp_rank( + State(state): State, + Json(body): Json, +) -> Response { + state.lock().unwrap().room_to_dp_rank.insert( + body.bootstrap_room.0, + RoomEntry { + dp_rank: body.dp_rank.0, + registered_at: Instant::now(), + }, + ); + "OK".into_response() +} + +#[derive(serde::Deserialize)] +struct QueryDpRanks { + bootstrap_rooms: Vec, +} + +/// Unknown rooms are silently omitted from the response, not an error. JSON +/// object keys are strings — Python's `str(room_int)` for free. +async fn query_dp_ranks(State(state): State, Json(body): Json) -> Response { + let reg = state.lock().unwrap(); + let result: HashMap = body + .bootstrap_rooms + .iter() + .filter_map(|room| { + let entry = reg.room_to_dp_rank.get(&room.0)?; + Some((room.0.to_string(), entry.dp_rank)) + }) + .collect(); + Json(result).into_response() +} + +/// No `/health` here: the merged api router already serves it (same 200 "OK" +/// the standalone Python bootstrap server answered, so probes are unchanged). +fn router(state: Shared) -> Router { + Router::new() + // Unmatched methods on a routed path get axum's built-in 405, matching + // Python's explicit method_not_allowed branch. + .route("/route", put(route_put).get(route_get)) + .route("/register_dp_rank", post(register_dp_rank)) + .route("/query_dp_ranks", post(query_dp_ranks)) + .with_state(state) +} + +/// Drop `room_to_dp_rank` entries older than `interval` — Python's +/// `_cleanup_expired_entries` loop (interval doubles as both period and TTL). +async fn cleanup_expired_entries(state: Shared, interval: Duration) { + loop { + tokio::time::sleep(interval).await; + state + .lock() + .unwrap() + .room_to_dp_rank + .retain(|_, entry| entry.registered_at.elapsed() <= interval); + } +} + +/// The registry routes plus their expiry sweeper, ready to mount on the api +/// router (`merge` the router, `tokio::spawn` the sweeper on the api runtime — +/// the runtime drop on shutdown cancels it along with the handlers). +pub(crate) fn router_and_sweeper() -> (Router, impl std::future::Future) { + let state = Shared::default(); + let cleanup_interval = Duration::from_secs(crate::environ::env_u64( + ENTRY_CLEANUP_INTERVAL_ENV, + ENTRY_CLEANUP_INTERVAL_DEFAULT_SECS, + )); + ( + router(state.clone()), + cleanup_expired_entries(state, cleanup_interval), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::runtime::{Runtime, RuntimeConfig, RustServerServerArgs, ServerArgs}; + use std::io::{Read, Write}; + use std::net::SocketAddr; + + /// Minimal HTTP/1.1 client (same style as the `runtime` tests): returns + /// `(status, body)`. + fn request( + addr: SocketAddr, + method: &str, + path_query: &str, + body: Option<&serde_json::Value>, + ) -> (u16, String) { + let body = body.map(|b| b.to_string()).unwrap_or_default(); + let mut conn = std::net::TcpStream::connect(addr).expect("connect"); + let req = format!( + "{method} {path_query} HTTP/1.1\r\nHost: t\r\nContent-Type: application/json\r\n\ + Content-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + conn.write_all(req.as_bytes()).unwrap(); + let mut response = String::new(); + conn.read_to_string(&mut response).unwrap(); + let status = response + .split_whitespace() + .nth(1) + .expect("status line") + .parse() + .expect("status code"); + let body = response + .split_once("\r\n\r\n") + .map(|(_, b)| b.to_string()) + .unwrap_or_default(); + (status, body) + } + + fn put_route(overrides: serde_json::Value) -> serde_json::Value { + let mut body = serde_json::json!({ + "attn_tp_size": 1, "attn_tp_rank": 0, + "attn_cp_size": 1, "attn_cp_rank": 0, + "attn_dp_size": 1, "attn_dp_rank": 0, + "pp_size": 1, "pp_rank": 0, + "system_dp_size": 1, "system_dp_rank": 0, + "rank_ip": "10.0.0.1", "rank_port": 17000, + "page_size": 64, "kv_cache_dtype": "auto", + "load_balance_method": "follow_bootstrap_room", + "enable_dsa_cache_layer_split": false, + "prefill_http_port": 30000, + }); + body.as_object_mut() + .unwrap() + .extend(overrides.as_object().unwrap().clone()); + body + } + + const SENTINEL: &str = + "/route?prefill_dp_rank=-1&prefill_cp_rank=-1&target_tp_rank=-1&target_pp_rank=-1"; + + /// Minimal prefill boot blob (same shape as the `runtime` tests): no + /// tokenizer load, the two mandatory `model_config` fields, and the + /// prefill role that mounts the registry. + const TEST_SERVER_ARGS: &str = r#"{ + "skip_tokenizer_init": true, + "disaggregation_mode": "prefill", + "model_config": {"context_len": 2048, "vocab_size": 1000} + }"#; + + /// Pick a free port (probe-bind pattern, as in the `runtime` tests) and + /// boot the full runtime there with the bootstrap registry mounted — the + /// registry serves on the api listener, so these tests also pin the merge + /// wiring (including the `enable_pd_bootstrap()` derivation from the + /// blob), not just the handlers. + fn start_on_free_port() -> (Runtime, SocketAddr) { + start_runtime(TEST_SERVER_ARGS) + } + + fn start_runtime(server_args_json: &str) -> (Runtime, SocketAddr) { + let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = probe.local_addr().unwrap(); + drop(probe); + let cfg = RuntimeConfig { + rust_server_args: RustServerServerArgs { + http_addr: addr, + api_worker_num: 1, + ..Default::default() + }, + server_args: Arc::new(ServerArgs::from_json(server_args_json).unwrap()), + }; + (crate::runtime::start(cfg).expect("start runtime"), addr) + } + + /// The full wire contract the Python decode side / PD router depends on: + /// 503 until every rank is registered, the `-1` sentinel returning + /// `PrefillServerInfo` with Python's exact field names, per-rank lookup + /// returning `PrefillRankInfo`, 404 for an unknown rank, 400 for missing + /// params, and `int(...)`-style acceptance of a string `rank_port`. Field + /// names and status codes are external literals owned by + /// disaggregation/common/conn.py — this pins the copy. + #[test] + fn route_contract_matches_python_client() { + let (_rt, addr) = start_on_free_port(); + + // Not registered yet → 503 (the decode side retries on exactly this). + let (status, _) = request(addr, "GET", SENTINEL, None); + assert_eq!(status, 503); + + // Missing/empty params → 400. + let (status, _) = request(addr, "GET", "/route?prefill_dp_rank=0", None); + assert_eq!(status, 400); + let (status, _) = request( + addr, + "GET", + "/route?prefill_dp_rank=0&prefill_cp_rank=&target_tp_rank=0&target_pp_rank=0", + None, + ); + assert_eq!(status, 400); + + // Register the single rank; rank_port as a STRING (Python coerces + // with `int(data["rank_port"])`, so the wire tolerates it). + let body = put_route(serde_json::json!({"rank_port": "17000"})); + let (status, text) = request(addr, "PUT", "/route", Some(&body)); + assert_eq!((status, text.as_str()), (200, "OK")); + + // Sentinel now serves the topology, keys verbatim from + // `dataclasses.asdict(PrefillServerInfo)`. + let (status, body) = request(addr, "GET", SENTINEL, None); + assert_eq!(status, 200); + let info: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!( + info, + serde_json::json!({ + "attn_tp_size": 1, "attn_cp_size": 1, "dp_size": 1, "pp_size": 1, + "page_size": 64, "kv_cache_dtype": "auto", + "follow_bootstrap_room": true, + "enable_dsa_cache_layer_split": false, + "prefill_http_port": 30000, + }) + ); + + // Per-rank lookup: `PrefillRankInfo` shape, rank_port back as an int. + let (status, body) = request( + addr, + "GET", + "/route?prefill_dp_rank=0&prefill_cp_rank=0&target_tp_rank=0&target_pp_rank=0", + None, + ); + assert_eq!(status, 200); + let rank: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!( + rank, + serde_json::json!({"rank_ip": "10.0.0.1", "rank_port": 17000}) + ); + + // Unknown rank → 404 (ready, but no such entry). + let (status, _) = request( + addr, + "GET", + "/route?prefill_dp_rank=1&prefill_cp_rank=0&target_tp_rank=0&target_pp_rank=0", + None, + ); + assert_eq!(status, 404); + } + + /// System-dp topology derivation: with `system_dp_size > 1` the dp axis + /// (readiness expectation AND rank keying) comes from `system_dp_*`, not + /// `attn_dp_*` — a "looks equivalent" simplification to always using + /// `attn_dp_*` passes single-dp tests but strands multi-dp deployments at + /// 503 / wrong-rank routes. + #[test] + fn system_dp_drives_readiness_and_rank_keys() { + let (_rt, addr) = start_on_free_port(); + + let rank0 = put_route(serde_json::json!({ + "system_dp_size": 2, "system_dp_rank": 0, "rank_ip": "10.0.0.1", + })); + let (status, _) = request(addr, "PUT", "/route", Some(&rank0)); + assert_eq!(status, 200); + + // dp_size resolved to system_dp_size=2 → one registration isn't ready. + let (status, _) = request(addr, "GET", SENTINEL, None); + assert_eq!(status, 503); + + let rank1 = put_route(serde_json::json!({ + "system_dp_size": 2, "system_dp_rank": 1, "rank_ip": "10.0.0.2", + })); + let (status, _) = request(addr, "PUT", "/route", Some(&rank1)); + assert_eq!(status, 200); + + let (status, body) = request(addr, "GET", SENTINEL, None); + assert_eq!(status, 200); + let info: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!(info["dp_size"], 2); + + // Ranks are keyed by system_dp_rank: dp=1 must be the second rank's ip. + let (status, body) = request( + addr, + "GET", + "/route?prefill_dp_rank=1&prefill_cp_rank=0&target_tp_rank=0&target_pp_rank=0", + None, + ); + assert_eq!(status, 200); + let rank: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!(rank["rank_ip"], "10.0.0.2"); + } + + /// The PD router's room→dp-rank side channel: register/query round-trip + /// with Python's `{str(room): dp_rank}` response shape, unknown rooms + /// silently omitted (not an error). (`/health` liveness now belongs to the + /// api router the registry is merged into.) + #[test] + fn dp_rank_round_trip() { + let (_rt, addr) = start_on_free_port(); + + let (status, text) = request( + addr, + "POST", + "/register_dp_rank", + Some(&serde_json::json!({"bootstrap_room": 42, "dp_rank": 3})), + ); + assert_eq!((status, text.as_str()), (200, "OK")); + + let (status, body) = request( + addr, + "POST", + "/query_dp_ranks", + Some(&serde_json::json!({"bootstrap_rooms": [42, 99]})), + ); + assert_eq!(status, 200); + let result: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!(result, serde_json::json!({"42": 3})); + } + + /// The registry mounts only on prefill (`enable_pd_bootstrap()`): a + /// non-prefill server must 404 the bootstrap routes rather than host an + /// empty replica — that replica would answer 503 "not registered" forever, + /// hiding a misdirected decode/router behind its retry loop. + #[test] + fn routes_absent_off_prefill() { + let non_prefill = r#"{ + "skip_tokenizer_init": true, + "model_config": {"context_len": 2048, "vocab_size": 1000} + }"#; + let (_rt, addr) = start_runtime(non_prefill); + + let (status, _) = request(addr, "GET", SENTINEL, None); + assert_eq!(status, 404); + let (status, _) = request( + addr, + "PUT", + "/route", + Some(&put_route(serde_json::json!({}))), + ); + assert_eq!(status, 404); + } +} diff --git a/rust/sglang-server/src/lib.rs b/rust/sglang-server/src/lib.rs index e82bf57a0..767a14c5c 100644 --- a/rust/sglang-server/src/lib.rs +++ b/rust/sglang-server/src/lib.rs @@ -61,7 +61,6 @@ impl Server { egress_ring_cap = 8192, channel_cap = 8192, cores = None, - server_args_json = "{}", ))] // pyo3 `#[new]` constructor: the wide arg list is the Python-facing boot diff --git a/rust/sglang-server/src/message/io_struct.rs b/rust/sglang-server/src/message/io_struct.rs index c3f0c7eb3..edc9b7f2e 100644 --- a/rust/sglang-server/src/message/io_struct.rs +++ b/rust/sglang-server/src/message/io_struct.rs @@ -29,7 +29,28 @@ wire_struct! { stream: bool, /// Not exposed by this server yet; the scheduler needs the slot filled. return_sampling_mask: bool, + return_flat_raw_top_logprobs: bool, return_hidden_states: bool, + /// Filler block (not exposed by this server yet): default/nil slots so + /// the PD fields below land on their Python wire indices (25–31). + return_routed_experts: bool, + routed_experts_start_len: i64, + return_indexer_topk: bool, + session_id: (), + session_params: (), + lora_id: (), + custom_logit_processor: (), + positional_embed_overrides: (), + /// PD-disaggregation block — the last fields emitted; everything after + /// `disagg_prefill_dp_rank` in Python has a msgspec default and is + /// omitted (short arrays decode with defaulted tails). + bootstrap_host: Option<&'a str>, + bootstrap_port: Option, + bootstrap_room: Option, + bootstrap_pair_key: Option<&'a str>, + decode_tp_size: Option, + routed_dp_rank: Option, + disagg_prefill_dp_rank: Option, } } @@ -74,7 +95,23 @@ impl<'a> From<&'a GenerateRequest> for TokenizedGenerateReqInput<'a> { token_ids_logprob: req.token_ids_logprob.as_ref(), stream: req.stream, return_sampling_mask: req.return_sampling_mask, + return_flat_raw_top_logprobs: false, return_hidden_states: req.return_hidden_states, + return_routed_experts: false, + routed_experts_start_len: 0, + return_indexer_topk: false, + session_id: (), + session_params: (), + lora_id: (), + custom_logit_processor: (), + positional_embed_overrides: (), + bootstrap_host: req.bootstrap_host.as_deref(), + bootstrap_port: req.bootstrap_port, + bootstrap_room: req.bootstrap_room, + bootstrap_pair_key: req.bootstrap_pair_key.as_deref(), + decode_tp_size: req.decode_tp_size, + routed_dp_rank: req.routed_dp_rank, + disagg_prefill_dp_rank: req.disagg_prefill_dp_rank, } } } @@ -141,12 +178,9 @@ mod tests { let bytes = TokenizedGenerateReqInput::from(&req).encode().unwrap(); let val = rmpv::decode::read_value(&mut &bytes[..]).unwrap(); let arr = val.as_array().expect("array"); - // msgspec requires >= 14 (through `stream`); we emit 16. - assert!( - arr.len() >= 14, - "header must have >=14 elements, got {}", - arr.len() - ); + // msgspec requires >= 14 (through `stream`); we emit 32 (through + // `disagg_prefill_dp_rank`). Trailing defaulted fields are omitted. + assert_eq!(arr.len(), 32, "header ends at disagg_prefill_dp_rank"); assert_eq!(arr[0].as_str(), Some("TokenizedGenerateReqInput")); assert_eq!(arr[1].as_str(), Some("r1")); assert!(arr[5].is_nil(), "idx 5 must be input_embeds (nil)"); @@ -166,8 +200,48 @@ mod tests { ); assert_eq!( arr[15].as_bool(), + Some(false), + "return_flat_raw_top_logprobs at idx 15" + ); + assert_eq!( + arr[16].as_bool(), Some(true), - "return_hidden_states at idx 15" + "return_hidden_states at idx 16" ); } + + /// The PD block must land on Python's wire indices 25–31, with the filler + /// block (17–24) holding its defaults — a shift here silently routes KV + /// transfers to the wrong host/room. + #[test] + fn header_bootstrap_block_is_positionally_aligned() { + let req = GenerateRequest { + rid: "r1".into(), + text: Some("hi".into()), + bootstrap_host: Some("10.0.0.1".into()), + bootstrap_port: Some(8998), + bootstrap_room: Some(i64::MAX), // routers draw from [0, 2^63) + bootstrap_pair_key: Some("pk".into()), + decode_tp_size: Some(2), + routed_dp_rank: Some(3), + disagg_prefill_dp_rank: Some(4), + ..Default::default() + }; + let bytes = TokenizedGenerateReqInput::from(&req).encode().unwrap(); + let val = rmpv::decode::read_value(&mut &bytes[..]).unwrap(); + let arr = val.as_array().expect("array"); + assert_eq!(arr[17].as_bool(), Some(false), "return_routed_experts"); + assert_eq!(arr[18].as_u64(), Some(0), "routed_experts_start_len"); + assert_eq!(arr[19].as_bool(), Some(false), "return_indexer_topk"); + for (i, slot) in arr.iter().enumerate().take(25).skip(20) { + assert!(slot.is_nil(), "idx {i} must be a nil default"); + } + assert_eq!(arr[25].as_str(), Some("10.0.0.1"), "bootstrap_host at 25"); + assert_eq!(arr[26].as_u64(), Some(8998), "bootstrap_port at 26"); + assert_eq!(arr[27].as_i64(), Some(i64::MAX), "bootstrap_room at 27"); + assert_eq!(arr[28].as_str(), Some("pk"), "bootstrap_pair_key at 28"); + assert_eq!(arr[29].as_i64(), Some(2), "decode_tp_size at 29"); + assert_eq!(arr[30].as_i64(), Some(3), "routed_dp_rank at 30"); + assert_eq!(arr[31].as_i64(), Some(4), "disagg_prefill_dp_rank at 31"); + } } diff --git a/rust/sglang-server/src/message/request.rs b/rust/sglang-server/src/message/request.rs index 00277a543..8c6fbf1eb 100644 --- a/rust/sglang-server/src/message/request.rs +++ b/rust/sglang-server/src/message/request.rs @@ -85,6 +85,27 @@ pub struct GenerateBody { /// Scalar-only in Python too (`return_text_in_logprobs: bool`). #[serde(default)] pub return_text_in_logprobs: Option, + // PD-disaggregation routing, injected per request by the PD router + // (mini_lb / sgl-model-gateway): a scalar for a single prompt, one-per-item + // lists for a batch. Elements are nullable (`List[Optional[...]]` in + // Python) — the router sends `bootstrap_port: [null, …]` when deferring to + // the scheduler's `--disaggregation-bootstrap-port` default. + #[serde(default)] + pub bootstrap_host: Option>>, + #[serde(default)] + pub bootstrap_port: Option>>, + /// `bootstrap_room` fits in i64: the PD routers draw it from `[0, 2^63)`. + #[serde(default)] + pub bootstrap_room: Option>>, + #[serde(default)] + pub bootstrap_pair_key: Option>>, + #[serde(default)] + pub decode_tp_size: Option>>, + /// DP routing hints — per-request scalars even for batches, as in Python. + #[serde(default)] + pub routed_dp_rank: Option, + #[serde(default)] + pub disagg_prefill_dp_rank: Option, } impl GenerateBody { @@ -107,6 +128,13 @@ impl GenerateBody { token_ids_logprob, return_hidden_states, return_text_in_logprobs, + bootstrap_host, + bootstrap_port, + bootstrap_room, + bootstrap_pair_key, + decode_tp_size, + routed_dp_rank, + disagg_prefill_dp_rank, // Unported `GenerateReqInput` fields land here and are dropped, as they // are on the Python path. .. @@ -271,6 +299,28 @@ impl GenerateBody { let top_logprobs_nums = fan_out(top_logprobs_num, n, "top_logprobs_num")?; let return_hidden = fan_out(return_hidden_states, n, "return_hidden_states")?; + // PD fields fan out like Python `_normalize_bootstrap_params`: scalars + // broadcast — except a scalar `bootstrap_room`, which becomes `room + i` + // (each item needs a distinct room; rooms are the P↔D pairing key). + // `fan_out` yields `Option>` for these nullable elements + // (outer: absent, inner: an explicit `null` element) — flatten, both + // mean "not set" downstream. + let bootstrap_hosts = flatten_column(fan_out(bootstrap_host, n, "bootstrap_host")?); + let bootstrap_ports = flatten_column(fan_out(bootstrap_port, n, "bootstrap_port")?); + let bootstrap_rooms = match bootstrap_room { + // `wrapping_add`, not `checked_`: rooms are drawn from `[0, 2^63)`, + // so a batch can only overflow by starting within `n` of `i64::MAX` + // — and distinct-but-wrapped still pairs P↔D, where saturating + // would collide every item onto one room. + Some(OneOrMany::One(Some(room))) => { + (0..n).map(|i| Some(room.wrapping_add(i as i64))).collect() + } + other => flatten_column(fan_out(other, n, "bootstrap_room")?), + }; + let bootstrap_pair_keys = + flatten_column(fan_out(bootstrap_pair_key, n, "bootstrap_pair_key")?); + let decode_tp_sizes = flatten_column(fan_out(decode_tp_size, n, "decode_tp_size")?); + // Every column above is exactly `n` long, so zip them by value: each // request takes ownership of its cell, with no indexing or bounds checks. let requests = izip!( @@ -283,6 +333,11 @@ impl GenerateBody { top_logprobs_nums, tid_logprobs, return_hidden, + bootstrap_hosts, + bootstrap_ports, + bootstrap_rooms, + bootstrap_pair_keys, + decode_tp_sizes, ) .map( |( @@ -295,6 +350,11 @@ impl GenerateBody { top_logprobs_num, token_ids_logprob, return_hidden_states, + bootstrap_host, + bootstrap_port, + bootstrap_room, + bootstrap_pair_key, + decode_tp_size, )| GenerateRequest { rid, text, @@ -311,6 +371,13 @@ impl GenerateBody { return_sampling_mask: false, // TODO: port Python's `return_sampling_mask` return_hidden_states: return_hidden_states.unwrap_or(false), return_text_in_logprobs, + bootstrap_host, + bootstrap_port, + bootstrap_room, + bootstrap_pair_key, + decode_tp_size, + routed_dp_rank, + disagg_prefill_dp_rank, }, ) .collect(); @@ -380,6 +447,18 @@ pub struct GenerateRequest { /// it is consumed on the way out, by `register_detok` → `DetokMsg::Register` /// → the shard's `decode_logprob_texts`. pub return_text_in_logprobs: Option, + /// PD-disaggregation routing, forwarded verbatim to the scheduler (which + /// fills a `None` port from `--disaggregation-bootstrap-port` and 400-aborts + /// a room-less request in PD mode). + pub bootstrap_host: Option, + pub bootstrap_port: Option, + pub bootstrap_room: Option, + pub bootstrap_pair_key: Option, + pub decode_tp_size: Option, + /// DP routing hints. The embedded server is rank-0-only (no DP controller), + /// so these are pure passthrough for the scheduler/LB protocol. + pub routed_dp_rank: Option, + pub disagg_prefill_dp_rank: Option, } impl GenerateRequest { @@ -439,6 +518,18 @@ impl HeapBytes for TokenIds { self.len() * std::mem::size_of::() } } +impl HeapBytes for Option { + fn heap_bytes(&self) -> usize { + self.as_ref().map_or(0, HeapBytes::heap_bytes) + } +} + +/// Collapse `fan_out`'s nullable-element output: outer `None` (field absent / +/// scalar broadcast of nothing) and inner `None` (an explicit `null` list +/// element) both mean "not set". +fn flatten_column(column: Vec>>) -> Vec> { + column.into_iter().map(Option::flatten).collect() +} /// Reject a broadcast whose clones would exceed [`MAX_BROADCAST_CLONE_BYTES`]. fn check_broadcast_budget(per_clone: usize, n: usize, name: &str) -> Result<(), Error> { @@ -821,4 +912,64 @@ mod tests { assert_eq!(a[0].rid.client_facing(), "same"); assert_eq!(b[0].rid.client_facing(), "same"); } + + /// PD bootstrap fields fan out like Python `_normalize_bootstrap_params`: + /// scalars broadcast, except a scalar `bootstrap_room` which becomes + /// `room + i` (each batch item needs a distinct room — rooms are the P↔D + /// pairing key); lists are per-item and must match the batch length. + #[test] + fn bootstrap_fields_fan_out() { + let (ps, _) = requests( + r#"{"text": ["a", "b"], "bootstrap_host": "h", "bootstrap_port": 8998, + "bootstrap_room": 7, "routed_dp_rank": 1}"#, + ) + .unwrap(); + for (i, p) in ps.iter().enumerate() { + assert_eq!(p.bootstrap_host.as_deref(), Some("h")); + assert_eq!(p.bootstrap_port, Some(8998)); + assert_eq!(p.bootstrap_room, Some(7 + i as i64)); + assert_eq!(p.routed_dp_rank, Some(1)); + } + + let (ps, _) = requests( + r#"{"text": ["a", "b"], "bootstrap_host": ["h1", "h2"], + "bootstrap_room": [10, 20]}"#, + ) + .unwrap(); + assert_eq!(ps[0].bootstrap_host.as_deref(), Some("h1")); + assert_eq!(ps[1].bootstrap_host.as_deref(), Some("h2")); + assert_eq!(ps[0].bootstrap_room, Some(10)); + assert_eq!(ps[1].bootstrap_room, Some(20)); + + let err = requests(r#"{"text": ["a", "b"], "bootstrap_room": [1, 2, 3]}"#).unwrap_err(); + assert!(err.to_string().contains("bootstrap_room"), "{err}"); + } + + /// The PD router (mini_lb) and PD-warmup payload shapes must parse. The + /// router sends `bootstrap_port: [null, …]` when no port was configured + /// (the scheduler fills its default) — null list elements must parse. + #[test] + fn accepts_pd_router_and_warmup_payloads() { + let (ps, _) = requests( + r#"{"text": ["a", "b"], "bootstrap_host": ["h", "h"], + "bootstrap_port": [null, null], + "bootstrap_room": [123456789, 987654321]}"#, + ) + .unwrap(); + assert_eq!(ps[0].bootstrap_host.as_deref(), Some("h")); + assert_eq!(ps[0].bootstrap_port, None); + assert_eq!(ps[1].bootstrap_room, Some(987654321)); + + let (ps, is_batch) = requests( + r#"{"sampling_params": {"temperature": 0.0, "max_new_tokens": 8, + "ignore_eos": true}, + "bootstrap_host": "2.2.2.2", "bootstrap_room": 0, + "input_ids": [10, 11, 12, 13], "routed_dp_rank": 0}"#, + ) + .unwrap(); + assert!(!is_batch); + assert_eq!(ps[0].bootstrap_host.as_deref(), Some("2.2.2.2")); + assert_eq!(ps[0].bootstrap_room, Some(0)); + assert_eq!(ps[0].routed_dp_rank, Some(0)); + } } diff --git a/rust/sglang-server/src/message/types.rs b/rust/sglang-server/src/message/types.rs index c15027555..360d8f5af 100644 --- a/rust/sglang-server/src/message/types.rs +++ b/rust/sglang-server/src/message/types.rs @@ -45,6 +45,12 @@ mod sealed { impl SealedItem for i64 {} impl SealedItem for String {} impl SealedItem for super::TokenIds {} + // Nullable elements for the PD bootstrap fields (`List[Optional[...]]` in + // Python — the PD router sends `bootstrap_port: [null, …]` when deferring to + // the scheduler's default port). A bare `null` never reaches `One(None)`: the + // outer `Option>` field consumes it first. + impl SealedItem for Option {} + impl SealedItem for Option {} } /// A msgspec `tag=True` struct: element 0 of its array is the Python class name diff --git a/rust/sglang-server/src/runtime/config.rs b/rust/sglang-server/src/runtime/config.rs index ca5a557a9..572acf391 100644 --- a/rust/sglang-server/src/runtime/config.rs +++ b/rust/sglang-server/src/runtime/config.rs @@ -96,6 +96,11 @@ pub struct ServerArgs { /// text. Matches the Python `TokenizerManager`. #[serde(default)] pub incremental_streaming_output: bool, + /// PD-disaggregation role: `"null"` (unified), `"prefill"`, or `"decode"`. + /// (On prefill, the KV bootstrap registry is mounted on the api router — + /// see [`Self::enable_pd_bootstrap`].) + #[serde(default = "default_disaggregation_mode")] + pub disaggregation_mode: String, /// The resolved Python `ModelConfig`, attached to the blob at dump time. #[serde(default)] pub model_config: ModelConfig, @@ -140,6 +145,14 @@ pub struct ModelConfig { pub vocab_size: Option, } +fn join_host_port(host: &str, port: u16) -> String { + if host.contains(':') && !host.starts_with('[') { + format!("[{host}]:{port}") // bare IPv6 (`::`) needs brackets to bind + } else { + format!("{host}:{port}") + } +} + fn default_host() -> String { "127.0.0.1".into() } @@ -149,6 +162,9 @@ fn default_port() -> u16 { fn default_log_level() -> String { "info".into() } +fn default_disaggregation_mode() -> String { + "null".into() +} fn default_worker_num() -> usize { 1 } @@ -170,13 +186,37 @@ impl ServerArgs { if self.model_config.vocab_size.is_none() { return Err("no resolvable vocab size (model_config.vocab_size)".into()); } + if !matches!( + self.disaggregation_mode.as_str(), + "null" | "prefill" | "decode" + ) { + return Err(format!( + "unknown disaggregation_mode '{}' in server_args", + self.disaggregation_mode + )); + } Ok(()) } + /// True on a prefill or decode node — requests need bootstrap routing. + pub fn is_disaggregation(&self) -> bool { + self.disaggregation_mode != "null" + } + + /// Serve the PD KV bootstrap registry on the api listener: every prefill + /// rust server hosts it, unconditionally — no extra topology gating. KV + /// managers and decode nodes reach the registry at the resolved + /// `disaggregation_bootstrap_port`, which rust-server mode aliases to the + /// api port, so whichever prefill server that port names is the one that + /// receives the registrations. + pub fn enable_pd_bootstrap(&self) -> bool { + self.disaggregation_mode == "prefill" + } + /// Bind address `host:port`. `host` is expected to be an IP — the result is - /// parsed as a `SocketAddr`. + /// parsed as a `SocketAddr`, so a bare IPv6 host gets bracketed. pub fn bind(&self) -> String { - format!("{}:{}", self.host, self.port) + join_host_port(&self.host, self.port) } /// Whether the HTTP access log is emitted, mirroring the Python server: diff --git a/rust/sglang-server/src/tokenizer.rs b/rust/sglang-server/src/tokenizer.rs index 2cc43bb26..0b553801b 100644 --- a/rust/sglang-server/src/tokenizer.rs +++ b/rust/sglang-server/src/tokenizer.rs @@ -73,13 +73,23 @@ pub fn resolve_model_file(path: &str, revision: Option<&str>, filename: &str) -> resolve_from_hub_cache(path, revision, filename) } -/// Locate a file for an HF Hub repo id in the local cache (`HF_HOME`). Offline — +/// Locate a file for an HF Hub repo id in the local cache. Offline — /// the scheduler pre-downloads the model. `None` if not cached. fn resolve_from_hub_cache(repo_id: &str, revision: Option<&str>, filename: &str) -> Option { use hf_hub::{Cache, Repo, RepoType}; + // Python resolves the cache dir as HF_HUB_CACHE > HUGGINGFACE_HUB_CACHE > + // HF_HOME/hub > ~/.cache/huggingface/hub; the hf-hub crate only knows + // HF_HOME. Honor the explicit cache-dir overrides first, or the Rust + // server misses models the Python scheduler already downloaded. + let cache = ["HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE"] + .iter() + .find_map(|var| std::env::var(var).ok()) + .map(|dir| Cache::new(dir.into())) + .unwrap_or_else(Cache::from_env); + let rev = revision.unwrap_or("main"); - Cache::from_env() + cache .repo(Repo::with_revision( repo_id.to_string(), RepoType::Model, diff --git a/test/registered/disaggregation/test_disaggregation_rust_server.py b/test/registered/disaggregation/test_disaggregation_rust_server.py new file mode 100644 index 000000000..86c9299b9 --- /dev/null +++ b/test/registered/disaggregation/test_disaggregation_rust_server.py @@ -0,0 +1,184 @@ +"""PD disaggregation with the embedded Rust server on both sides. + +Same 2-GPU layout as test_disaggregation_basic (prefill GPU 0, decode GPU 1, +mini_lb in front), but prefill and decode run with ``SGLANG_RUST_SERVER=1`` — +covering the Rust `/generate` bootstrap-field intake (scalar form via the gsm8k +eval's single-prompt requests, per-item list form via the batch test), the +positional scheduler-wire PD block, the KV bootstrap registry served on the +rust api listener, the PD warmup fan-out, and the fake-bootstrap health probe. + +The Rust server has no OpenAI endpoints, so everything (including the gsm8k +eval) goes through ``/generate``. + +Usage: +python3 -m unittest test_disaggregation_rust_server.TestDisaggregationRustServer +""" + +import json +import unittest +from types import SimpleNamespace + +import requests + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_eval import run_eval +from sglang.test.server_fixtures.disaggregation_fixture import ( + PDDisaggregationServerBase, +) +from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST, is_rust_server_built + +register_cuda_ci(est_time=500, stage="base-b", runner_config="2-gpu-large") + + +@unittest.skipUnless( + is_rust_server_built(), + "embedded rust server extension not built", +) +class TestDisaggregationRustServer(PDDisaggregationServerBase): + extra_prefill_env = {"SGLANG_RUST_SERVER": "1"} + extra_decode_env = {"SGLANG_RUST_SERVER": "1"} + + @classmethod + def setUpClass(cls): + super().setUpClass() + # Rust-server prefill serves the KV bootstrap registry on its api + # listener (a separate --disaggregation-bootstrap-port is a launch + # error there), so point both sides' bootstrap port at it: decode's + # flag is its fallback for requests without a bootstrap_port field, + # which is what mini_lb sends when --prefill carries no port. + cls.bootstrap_port = cls.prefill_port + cls.model = DEFAULT_MODEL_NAME_FOR_TEST + # launch_all already exercises the PD-specific plumbing: the rust PD + # warmup fan-out and the fake-bootstrap /health probe on both sides. + cls.launch_all() + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.lb_url, + eval_name="gsm8k", + api="generate", # the Rust server has no /v1/completions + max_tokens=512, + num_examples=64, + num_threads=32, + ) + metrics = run_eval(args) + print(f"Evaluation metrics: {metrics}") + self.assertGreater(metrics["score"], 0.62) + + def test_generate_stream_via_lb(self): + # The scalar-bootstrap non-stream path is already covered 64x with an + # accuracy gate by test_gsm8k; what is unique here is mini_lb passing + # the decode node's SSE frames through under PD. + response = requests.post( + self.lb_url + "/generate", + json={ + "text": "The capital of France is", + "sampling_params": {"temperature": 0, "max_new_tokens": 16}, + "stream": True, + }, + stream=True, + ) + self.assertEqual(response.status_code, 200) + chunks = [] + for line in response.iter_lines(decode_unicode=True): + if not line or not line.startswith("data:"): + continue + payload = line[len("data:") :].strip() + if payload == "[DONE]": + break + chunks.append(json.loads(payload)) + self.assertTrue(chunks) + self.assertTrue(chunks[-1]["text"]) + # Frames are cumulative (--incremental-streaming-output defaults off), + # so the last frame must extend the first. (Frame *count* is not + # asserted: a slow reader legitimately coalesces a drained backlog.) + self.assertTrue(chunks[-1]["text"].startswith(chunks[0]["text"])) + # Exactly one terminal frame, and it is the last one. On a PD stream the + # prefill node produces its own finish_reason frame; leaking that into + # the decode stream would truncate the client mid-generation. + terminal = [ + i + for i, chunk in enumerate(chunks) + if chunk["meta_info"]["finish_reason"] is not None + ] + self.assertEqual(terminal, [len(chunks) - 1], f"{terminal=} {len(chunks)=}") + # One request id across the whole stream — not prefill's, then decode's. + self.assertEqual(len({chunk["meta_info"]["id"] for chunk in chunks}), 1) + + def test_batch_generate_via_lb(self): + # A batch makes the router inject per-item bootstrap lists — the list + # intake + per-item fan-out path on the Rust side. + response = requests.post( + self.lb_url + "/generate", + json={ + "text": ["The capital of France is", "The capital of Japan is"], + "sampling_params": {"temperature": 0, "max_new_tokens": 16}, + }, + ) + self.assertEqual(response.status_code, 200) + j = response.json() + self.assertEqual(len(j), 2) + # Per-prompt answers, not just non-empty text: a bootstrap room paired + # with the wrong list index hands one item the other's transferred KV, + # which a truthiness check cannot see. + for item, expected in zip(j, ("paris", "tokyo")): + self.assertIn(expected, item["text"].lower(), item) + self.assertIsNotNone(item["meta_info"]["finish_reason"]) + + def test_logprob_merge_via_lb(self): + # With return_logprob the router merges the *prefill* response's + # input_token_logprobs into the decode response — both sides must + # produce complete logprob meta_info. (No `return_input_logprob` here: + # the Rust /generate body does not declare it.) + response = requests.post( + self.lb_url + "/generate", + json={ + "text": "The capital of France is", + "sampling_params": {"temperature": 0, "max_new_tokens": 16}, + "return_logprob": True, + "logprob_start_len": 0, + }, + ) + self.assertEqual(response.status_code, 200) + meta = response.json()["meta_info"] + self.assertEqual(len(meta["output_token_logprobs"]), meta["completion_tokens"]) + # The *whole* prompt, since logprob_start_len is 0: a merge that drops + # prefill's list and leaves only what decode itself saw still yields a + # non-empty list, so pin the exact length. + self.assertEqual(len(meta["input_token_logprobs"]), meta["prompt_tokens"]) + + def test_missing_bootstrap_is_rejected(self): + # Negative branch of the fake-bootstrap health probe: a /generate that + # reaches a PD node *without* the router's bootstrap fields must surface + # the scheduler's 400 abort through the rust wire — not hang, not 500. + # Nothing else in this suite reaches the rust egress' abort_status path. + response = requests.post( + self.prefill_url + "/generate", + json={ + "text": "The capital of France is", + "sampling_params": {"temperature": 0, "max_new_tokens": 16}, + }, + timeout=60, + ) + self.assertEqual(response.status_code, 400, response.text) + + def test_backend_health(self): + # /health_generate directly on each side: on a PD node the probe only + # passes with the fake bootstrap pair injected (room-less requests are + # 400-aborted by the scheduler). Not the fixture's assert_process_healthy: + # its 10s client timeout is shorter than the probe's own deadline + # (SGLANG_HEALTH_CHECK_TIMEOUT, 20s), which would turn a slow-but-passing + # side into a connection error. + for name, process, url in ( + ("prefill", self.process_prefill, self.prefill_url), + ("decode", self.process_decode, self.decode_url), + ): + self.assertIsNone( + process.poll(), f"{name} exited with code {process.returncode}" + ) + response = requests.get(url + "/health_generate", timeout=60) + self.assertEqual(response.status_code, 200, response.text) + + +if __name__ == "__main__": + unittest.main()