[Router] Add a rust-based router (#1790)

rust/src/lib.rs

@@ -0,0 +1,45 @@
+use pyo3::prelude::*;
+mod server;
+pub mod router;
+
+// Python binding
+#[pyclass]
+struct Router {
+    host: String,
+    port: u16,
+    worker_urls: Vec<String>,
+    policy: String
+}
+
+#[pymethods]
+impl Router {
+    #[new]
+    fn new(host: String, port: u16, worker_urls: Vec<String>, policy: String) -> Self {
+        Router {
+            host,
+            port,
+            worker_urls,
+            policy
+        }
+    }
+
+    fn start(&self) -> PyResult<()> {
+        let host = self.host.clone();
+        let port = self.port;
+        let worker_urls = self.worker_urls.clone();
+        let policy = self.policy.clone();
+
+        actix_web::rt::System::new().block_on(async move {
+            server::startup(host, port, worker_urls, policy).await.unwrap();
+        });
+
+        Ok(())
+    }
+}
+
+// python usage: `from sglang_router import Router`
+#[pymodule]
+fn sglang_router(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<Router>()?;
+    Ok(())
+}

rust/src/main.rs

@@ -0,0 +1,28 @@
+// src/main.rs
+use clap::Parser;
+use clap::builder::PossibleValuesParser;
+// declare child modules
+mod server;
+mod router;
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    #[arg(long, default_value = "127.0.0.1")]
+    host: String,
+
+    #[arg(long, default_value_t = 3001)]
+    port: u16,
+
+    #[arg(long, value_delimiter = ',')]
+    worker_urls: Vec<String>,
+
+    #[arg(long, default_value = "round_robin", value_parser = PossibleValuesParser::new(&["round_robin", "random"]))]
+    policy: String,
+}
+
+#[actix_web::main]
+async fn main() -> std::io::Result<()> {
+    let args = Args::parse();
+    server::startup(args.host, args.port, args.worker_urls, args.policy).await
+}

rust/src/router.rs

@@ -0,0 +1,88 @@
+// src/router.rs
+
+use std::fmt::Debug;
+
+/// Generic Router trait that can be implemented with different policies
+pub trait Router: Send + Sync + Debug {
+    /// Select a worker URL based on the implementation's policy
+    /// Returns None if no worker is available
+    fn select(&self) -> Option<String>;
+
+    // get first worker
+    fn get_first(&self) -> Option<String>;
+}
+
+// Round Robin Router
+#[derive(Debug)]
+pub struct RoundRobinRouter {
+    worker_urls: Vec<String>,
+    current_index: std::sync::atomic::AtomicUsize, // AtomicUsize is a thread-safe integer
+}
+
+impl RoundRobinRouter {
+    pub fn new(worker_urls: Vec<String>) -> Self {
+        Self {
+            worker_urls,
+            current_index: std::sync::atomic::AtomicUsize::new(0),
+        }
+    }
+}
+
+impl Router for RoundRobinRouter {
+    fn select(&self) -> Option<String> {
+        if self.worker_urls.is_empty() {
+            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) 
+            % self.worker_urls.len();
+        Some(self.worker_urls[index].clone())
+    }
+
+    fn get_first(&self) -> Option<String> {
+        if self.worker_urls.is_empty() {
+            return None;
+        }
+        Some(self.worker_urls[0].clone())
+    }
+}
+
+// Random Router
+#[derive(Debug)]
+pub struct RandomRouter {
+    worker_urls: Vec<String>,
+}
+
+impl RandomRouter {
+    pub fn new(worker_urls: Vec<String>) -> Self {
+        Self { worker_urls }
+    }
+}
+
+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()
+    }
+
+    fn get_first(&self) -> Option<String> {
+        if self.worker_urls.is_empty() {
+            return None;
+        }
+        Some(self.worker_urls[0].clone())
+    }
+}
+
+// create a router based on routing policy
+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),
+    }
+}

rust/src/server.rs

@@ -0,0 +1,168 @@
+use actix_web::{get, post, web, App, HttpServer, HttpResponse, HttpRequest, 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 {
+    router: Box<dyn Router>,
+    client: reqwest::Client,
+}
+
+
+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,
+        }
+    }
+}
+
+#[get("/v1/models")]
+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))
+        .send()
+        .await 
+    {
+        Ok(res) => {
+            let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
+            .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() {
+        Some(url) => url,
+        None => return HttpResponse::InternalServerError().finish(),
+    };
+    // Use the shared client
+    match data.client
+        .get(&format!("{}/get_model_info", worker_url))
+        .send()
+        .await 
+    {
+        Ok(res) => {
+            let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
+            .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 {
+
+    // create a router struct
+    // TODO: use router abstraction for different policy
+    let worker_url= match data.router.select() {
+        Some(url) => url,
+        None => return HttpResponse::InternalServerError().finish(),
+    };
+
+    // Check if client requested streaming
+    let is_stream = serde_json::from_slice::<serde_json::Value>(&body)
+        .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))
+        .header(
+            "Content-Type", 
+            req.headers()
+                .get("Content-Type")
+                .and_then(|h| h.to_str().ok())
+                .unwrap_or("application/json")
+        )
+        .body(body.to_vec())
+        .send()
+        .await 
+    {
+        Ok(res) => res,
+        Err(_) => return HttpResponse::InternalServerError().finish(),
+    };
+
+    let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
+        .unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
+
+    if !is_stream {
+        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"))),
+            }))
+    }
+}
+
+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);
+
+    // Create client once with configuration
+    let client = reqwest::Client::builder()
+        .build()
+        .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,
+    ));
+
+    HttpServer::new(move || {
+        App::new()
+            .app_data(app_state.clone())
+            .service(generate)
+            .service(v1_model)
+            .service(get_model_info)
+    })
+    .bind((host, port))?
+    .run()
+    .await
+}