[rust-server] PD disaggregation support (#33125)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Rain Jiang <rain-jiang@outlook.com>
This commit is contained in:
Kan Wu
co-authored by
Claude Fable 5
Cursor
Rain Jiang
Claude Fable 5
Cursor
Rain Jiang
parent
c844244da5
commit
5d2dbb35a6
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -61,6 +61,10 @@ fn health_routes() -> Router<AppState> {
|
||||
.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<AppState>, 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_<uuid>` rid form
|
||||
rid: Rid::new_health_check(),
|
||||
@@ -91,6 +99,8 @@ async fn health_generate(State(state): State<AppState>, 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) =
|
||||
|
||||
@@ -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<i64>,
|
||||
kv_cache_dtype: Option<String>,
|
||||
follow_bootstrap_room: bool,
|
||||
enable_dsa_cache_layer_split: bool,
|
||||
prefill_http_port: Option<i64>,
|
||||
}
|
||||
|
||||
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<i64>,
|
||||
attn_cp_size: Option<i64>,
|
||||
dp_size: Option<i64>,
|
||||
pp_size: Option<i64>,
|
||||
page_size: Option<i64>,
|
||||
kv_cache_dtype: Option<String>,
|
||||
follow_bootstrap_room: Option<bool>,
|
||||
enable_dsa_cache_layer_split: Option<bool>,
|
||||
prefill_http_port: Option<i64>,
|
||||
/// 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<i64, RoomEntry>,
|
||||
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<i64> {
|
||||
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<Mutex<Registry>>;
|
||||
|
||||
/// 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: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
|
||||
#[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<String>,
|
||||
#[serde(default)]
|
||||
prefill_http_port: Option<Int>,
|
||||
#[serde(default)]
|
||||
load_balance_method: Option<String>,
|
||||
#[serde(default)]
|
||||
enable_dsa_cache_layer_split: Option<bool>,
|
||||
}
|
||||
|
||||
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<Shared>, Json(body): Json<RoutePut>) -> 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<Shared>,
|
||||
Query(query): Query<HashMap<String, String>>,
|
||||
) -> Response {
|
||||
// A missing, empty (Python truthiness), or non-integer param → 400.
|
||||
let rank = |k: &str| query.get(k).and_then(|v| v.trim().parse::<i64>().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<Shared>,
|
||||
Json(body): Json<RegisterDpRank>,
|
||||
) -> 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<Int>,
|
||||
}
|
||||
|
||||
/// 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<Shared>, Json(body): Json<QueryDpRanks>) -> Response {
|
||||
let reg = state.lock().unwrap();
|
||||
let result: HashMap<String, i64> = 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<Output = ()>) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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<i64>,
|
||||
bootstrap_room: Option<i64>,
|
||||
bootstrap_pair_key: Option<&'a str>,
|
||||
decode_tp_size: Option<i64>,
|
||||
routed_dp_rank: Option<i64>,
|
||||
disagg_prefill_dp_rank: Option<i64>,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<bool>,
|
||||
// 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<OneOrMany<Option<String>>>,
|
||||
#[serde(default)]
|
||||
pub bootstrap_port: Option<OneOrMany<Option<i64>>>,
|
||||
/// `bootstrap_room` fits in i64: the PD routers draw it from `[0, 2^63)`.
|
||||
#[serde(default)]
|
||||
pub bootstrap_room: Option<OneOrMany<Option<i64>>>,
|
||||
#[serde(default)]
|
||||
pub bootstrap_pair_key: Option<OneOrMany<Option<String>>>,
|
||||
#[serde(default)]
|
||||
pub decode_tp_size: Option<OneOrMany<Option<i64>>>,
|
||||
/// DP routing hints — per-request scalars even for batches, as in Python.
|
||||
#[serde(default)]
|
||||
pub routed_dp_rank: Option<i64>,
|
||||
#[serde(default)]
|
||||
pub disagg_prefill_dp_rank: Option<i64>,
|
||||
}
|
||||
|
||||
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<Option<T>>` 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<bool>,
|
||||
/// 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<String>,
|
||||
pub bootstrap_port: Option<i64>,
|
||||
pub bootstrap_room: Option<i64>,
|
||||
pub bootstrap_pair_key: Option<String>,
|
||||
pub decode_tp_size: Option<i64>,
|
||||
/// 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<i64>,
|
||||
pub disagg_prefill_dp_rank: Option<i64>,
|
||||
}
|
||||
|
||||
impl GenerateRequest {
|
||||
@@ -439,6 +518,18 @@ impl HeapBytes for TokenIds {
|
||||
self.len() * std::mem::size_of::<i32>()
|
||||
}
|
||||
}
|
||||
impl<T: HeapBytes> HeapBytes for Option<T> {
|
||||
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<T>(column: Vec<Option<Option<T>>>) -> Vec<Option<T>> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<OneOrMany<…>>` field consumes it first.
|
||||
impl SealedItem for Option<i64> {}
|
||||
impl SealedItem for Option<String> {}
|
||||
}
|
||||
|
||||
/// A msgspec `tag=True` struct: element 0 of its array is the Python class name
|
||||
|
||||
@@ -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<u64>,
|
||||
}
|
||||
|
||||
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:
|
||||
|
||||
@@ -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<String> {
|
||||
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,
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user