EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[router] Impl radix tree and set up CI (#1893)

Browse Source 
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
This commit is contained in:
Byron Hsu
2024-11-04 10:56:52 -08:00
committed by GitHub
parent 3cd2809277
commit 530ff541cf
11 changed files with 458 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
rust/src/lib.rs
View File

@@ -1,6 +1,7 @@
use pyo3::prelude::*;
mod server;
pub mod router;
mod server;
pub mod tree;
// Python binding
#[pyclass]
@@ -8,7 +9,7 @@ struct Router {
host: String,
port: u16,
worker_urls: Vec<String>,
policy: String
policy: String,
}
#[pymethods]
@@ -19,7 +20,7 @@ impl Router {
host,
port,
worker_urls,
policy
policy,
}
}
@@ -30,7 +31,9 @@ impl Router {
let policy = self.policy.clone();
actix_web::rt::System::new().block_on(async move {
server::startup(host, port, worker_urls, policy).await.unwrap();
server::startup(host, port, worker_urls, policy)
.await
.unwrap();
});
Ok(())
@@ -42,4 +45,4 @@ impl Router {
fn sglang_router(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<Router>()?;
Ok(())
}
}

7
rust/src/main.rs
View File

@@ -1,9 +1,10 @@
// src/main.rs
use clap::Parser;
use clap::builder::PossibleValuesParser;
use clap::Parser;
// declare child modules
mod server;
mod router;
mod server;
mod tree;
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
@@ -25,4 +26,4 @@ struct Args {
async fn main() -> std::io::Result<()> {
let args = Args::parse();
server::startup(args.host, args.port, args.worker_urls, args.policy).await
}
}

15
rust/src/router.rs
View File

@@ -34,7 +34,9 @@ impl Router for RoundRobinRouter {
return None;
}
// Use relaxed because operation order doesn't matter in round robin
let index = self.current_index.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
let index = self
.current_index
.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
% self.worker_urls.len();
Some(self.worker_urls[index].clone())
}
@@ -62,11 +64,11 @@ impl RandomRouter {
impl Router for RandomRouter {
fn select(&self) -> Option<String> {
use rand::seq::SliceRandom;
if self.worker_urls.is_empty() {
return None;
}
self.worker_urls.choose(&mut rand::thread_rng()).cloned()
}
@@ -83,6 +85,9 @@ pub fn create_router(worker_urls: Vec<String>, policy: String) -> Box<dyn Router
match policy.to_lowercase().as_str() {
"random" => Box::new(RandomRouter::new(worker_urls)),
"round_robin" => Box::new(RoundRobinRouter::new(worker_urls)),
_ => panic!("Unknown routing policy: {}. The available policies are 'random' and 'round_robin'", policy),
_ => panic!(
"Unknown routing policy: {}. The available policies are 'random' and 'round_robin'",
policy
),
}
}
}

99
rust/src/server.rs
View File

@@ -1,10 +1,9 @@
use actix_web::{get, post, web, App, HttpServer, HttpResponse, HttpRequest, Responder};
use crate::router::create_router;
use crate::router::Router;
use actix_web::http::header::{HeaderValue, CONTENT_TYPE};
use actix_web::{get, post, web, App, HttpRequest, HttpResponse, HttpServer, Responder};
use bytes::Bytes;
use futures_util::StreamExt;
use actix_web::http::header::{HeaderValue, CONTENT_TYPE};
use crate::router::Router;
use crate::router::create_router;
#[derive(Debug)]
pub struct AppState {
@@ -12,89 +11,77 @@ pub struct AppState {
client: reqwest::Client,
}
impl AppState
{
impl AppState {
pub fn new(worker_urls: Vec<String>, policy: String, client: reqwest::Client) -> Self {
// Create router based on policy
let router = create_router(worker_urls, policy);
Self {
router,
client,
}
Self { router, client }
}
}
#[get("/v1/models")]
async fn v1_model(
data: web::Data<AppState>,
) -> impl Responder {
let worker_url= match data.router.get_first() {
async fn v1_model(data: web::Data<AppState>) -> impl Responder {
let worker_url = match data.router.get_first() {
Some(url) => url,
None => return HttpResponse::InternalServerError().finish(),
};
// Use the shared client
match data.client
.get(&format!("{}/v1/models", worker_url))
match data
.client
.get(format!("{}/v1/models", worker_url))
.send()
.await
.await
{
Ok(res) => {
let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
.unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
.unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
// print the status
println!("Worker URL: {}, Status: {}", worker_url, status);
match res.bytes().await {
Ok(body) => HttpResponse::build(status).body(body.to_vec()),
Err(_) => HttpResponse::InternalServerError().finish(),
}
},
}
Err(_) => HttpResponse::InternalServerError().finish(),
}
}
#[get("/get_model_info")]
async fn get_model_info(
data: web::Data<AppState>,
) -> impl Responder {
let worker_url= match data.router.get_first() {
async fn get_model_info(data: web::Data<AppState>) -> impl Responder {
let worker_url = match data.router.get_first() {
Some(url) => url,
None => return HttpResponse::InternalServerError().finish(),
};
// Use the shared client
match data.client
.get(&format!("{}/get_model_info", worker_url))
match data
.client
.get(format!("{}/get_model_info", worker_url))
.send()
.await
.await
{
Ok(res) => {
let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
.unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
.unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
// print the status
println!("Worker URL: {}, Status: {}", worker_url, status);
match res.bytes().await {
Ok(body) => HttpResponse::build(status).body(body.to_vec()),
Err(_) => HttpResponse::InternalServerError().finish(),
}
},
}
Err(_) => HttpResponse::InternalServerError().finish(),
}
}
// no deser and ser, just forward and return
#[post("/generate")]
async fn generate(
req: HttpRequest,
body: Bytes,
data: web::Data<AppState>,
) -> impl Responder {
async fn generate(req: HttpRequest, body: Bytes, data: web::Data<AppState>) -> impl Responder {
// create a router struct
// TODO: use router abstraction for different policy
let worker_url= match data.router.select() {
let worker_url = match data.router.select() {
Some(url) => url,
None => return HttpResponse::InternalServerError().finish(),
};
@@ -104,18 +91,19 @@ async fn generate(
.map(|v| v.get("stream").and_then(|s| s.as_bool()).unwrap_or(false))
.unwrap_or(false);
let res = match data.client
.post(&format!("{}/generate", worker_url))
let res = match data
.client
.post(format!("{}/generate", worker_url))
.header(
"Content-Type",
"Content-Type",
req.headers()
.get("Content-Type")
.and_then(|h| h.to_str().ok())
.unwrap_or("application/json")
.unwrap_or("application/json"),
)
.body(body.to_vec())
.send()
.await
.await
{
Ok(res) => res,
Err(_) => return HttpResponse::InternalServerError().finish(),
@@ -128,18 +116,25 @@ async fn generate(
match res.bytes().await {
Ok(body) => HttpResponse::build(status).body(body.to_vec()),
Err(_) => HttpResponse::InternalServerError().finish(),
}
}
} else {
HttpResponse::build(status)
.insert_header((CONTENT_TYPE, HeaderValue::from_static("text/event-stream")))
.streaming(res.bytes_stream().map(|b| match b {
Ok(b) => Ok::<_, actix_web::Error>(b),
Err(_) => Err(actix_web::Error::from(actix_web::error::ErrorInternalServerError("Failed to read stream"))),
Err(_) => Err(actix_web::error::ErrorInternalServerError(
"Failed to read stream",
)),
}))
}
}
pub async fn startup(host: String, port: u16, worker_urls: Vec<String>, routing_policy: String) -> std::io::Result<()> {
pub async fn startup(
host: String,
port: u16,
worker_urls: Vec<String>,
routing_policy: String,
) -> std::io::Result<()> {
println!("Starting server on {}:{}", host, port);
println!("Worker URLs: {:?}", worker_urls);
@@ -149,11 +144,7 @@ pub async fn startup(host: String, port: u16, worker_urls: Vec<String>, routing_
.expect("Failed to create HTTP client");
// Store both worker_urls and client in AppState
let app_state = web::Data::new(AppState::new(
worker_urls,
routing_policy,
client,
));
let app_state = web::Data::new(AppState::new(worker_urls, routing_policy, client));
HttpServer::new(move || {
App::new()
@@ -165,4 +156,4 @@ pub async fn startup(host: String, port: u16, worker_urls: Vec<String>, routing_
.bind((host, port))?
.run()
.await
}
}

184
rust/src/tree.rs Normal file
View File

@@ -0,0 +1,184 @@
use std::collections::HashMap;
use std::mem;
#[derive(Clone)]
pub struct Node {
pub children: HashMap<usize, Node>, // the key is first id of the child because each child must have unique first id
pub ids: Vec<usize>,
pub count: usize,
}
pub struct RadixTree {
pub root: Node,
}
fn common_prefix_len(a: &[usize], b: &[usize]) -> usize {
let mut i = 0;
while i < a.len() && i < b.len() && a[i] == b[i] {
i += 1;
}
i
}
impl Default for RadixTree {
fn default() -> Self {
Self::new()
}
}
impl RadixTree {
pub fn new() -> Self {
RadixTree {
root: Node {
children: HashMap::new(),
ids: Vec::new(),
count: 0,
},
}
}
pub fn insert(&mut self, input_ids: &[usize]) {
let mut curr = &mut self.root;
curr.count += 1;
let mut curr_idx = 0;
let input_ids_len = input_ids.len();
while curr_idx < input_ids_len {
let first_id = &input_ids[curr_idx];
// TODO: changing this get_mut causes error
if curr.children.contains_key(first_id) {
let child = curr.children.get_mut(first_id).unwrap();
let prefix_len = common_prefix_len(&input_ids[curr_idx..], &child.ids);
if prefix_len == child.ids.len() {
// move curr to child
curr = child;
curr.count += 1;
curr_idx += prefix_len;
} else {
// split child
// [child]->... => [child]->[new child]->...
let new_child = Node {
// to avoid clone: replace child.children with default value (empty vector) and return the original value
children: mem::take(&mut child.children),
ids: child.ids[prefix_len..].to_vec(),
count: child.count,
};
child.ids = child.ids[..prefix_len].to_vec();
child.children = HashMap::new();
child.children.insert(new_child.ids[0], new_child);
curr = child;
curr.count += 1;
curr_idx += prefix_len;
}
} else {
// create new child
let new_child = Node {
children: HashMap::new(),
ids: input_ids[curr_idx..].to_vec(),
count: 0,
};
let first_id = new_child.ids[0];
curr.children.insert(first_id, new_child);
curr = curr.children.get_mut(&first_id).unwrap();
curr.count += 1;
curr_idx = input_ids_len;
}
}
}
pub fn prefix_match<'a>(&self, input_ids: &'a [usize]) -> &'a [usize] {
let mut curr = &self.root;
let mut curr_idx = 0;
let input_ids_len = input_ids.len();
while curr_idx < input_ids_len {
match curr.children.get(&input_ids[curr_idx]) {
Some(child) => {
let prefix_len = common_prefix_len(&input_ids[curr_idx..], &child.ids);
if prefix_len == child.ids.len() {
curr_idx += prefix_len;
curr = child;
} else {
curr_idx += prefix_len;
break;
}
}
None => {
break;
}
}
}
&input_ids[..curr_idx]
}
pub fn delete(&mut self, input_ids: &[usize]) {
let mut curr = &mut self.root;
curr.count -= 1;
let mut curr_idx = 0;
let input_ids_len = input_ids.len();
while curr_idx < input_ids_len {
let first_id = &input_ids[curr_idx];
if curr.children.contains_key(first_id) {
let child = curr.children.get(first_id).unwrap();
let prefix_len = common_prefix_len(&input_ids[curr_idx..], &child.ids);
if prefix_len == child.ids.len() {
if child.count == 1 {
// If count will become 0, remove the child
let child = curr.children.get_mut(first_id).unwrap();
child.count -= 1;
curr.children.remove(first_id);
break;
} else {
// Otherwise decrement count and continue
let child = curr.children.get_mut(first_id).unwrap();
child.count -= 1;
curr = child;
curr_idx += prefix_len;
}
} else {
panic!("No match found for {:?}", input_ids);
}
} else {
panic!("No match found for {:?}", input_ids);
}
}
}
// for debug
pub fn pretty_print(&self) {
println!("RadixTree:");
Self::print_node(&self.root, String::from(""));
}
fn print_node(node: &Node, prefix: String) {
// Print current node info with "count" word
println!("{}└── {:?} (count: {})", prefix, node.ids, node.count);
// Print children with proper prefixes
for (i, child) in node.children.values().enumerate() {
let is_last = i == node.children.len() - 1;
let child_prefix = if is_last {
format!("{} ", prefix) // Add space for last child
} else {
format!("{}│ ", prefix) // Add vertical line for other children
};
Self::print_node(child, child_prefix);
}
}
}