[router] rust-based router

rust/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "sglang-router"
+version = "0.0.0"
+edition = "2021"
+
+[[bin]]
+name = "router"
+path = "src/main.rs"
+
+[lib]
+name = "router"
+crate-type = ["cdylib"]
+
+[dependencies]
+actix-web = "4.0"
+serde = { version = "1.0", features = ["derive"] }
+clap = { version = "4.4", features = ["derive"] }
+bytes = "1.8.0"
+rand = "0.8.5"
+reqwest = { version = "0.12.8", features = ["stream"] }
+futures-util = "0.3"
+serde_json = "=1.0.1"
+pyo3 = { version = "0.22.5", features = ["extension-module"] }

rust/readme.md

@@ -0,0 +1,77 @@
+# SGLang Router
+
+SGLang router is a standalone module implemented in Rust to achieve data parallelism across SGLang instances.
+
+### Installation
+
+1. Install Rust
+```bash
+# Install rustup (Rust installer and version manager)
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
+
+# Follow the installation prompts, then reload your shell
+source $HOME/.cargo/env
+
+# Verify installation
+rustc --version
+cargo --version
+```
+
+2. Build the router
+```bash
+# Navigate to the rust directory
+cd ./rust
+
+# Build the project
+cargo build
+
+# Verify the binary works correctly
+./target/debug/router --help
+```
+
+The help command will show available options:
+```
+Usage: router [OPTIONS]
+
+Options:
+      --host <HOST>                [default: 127.0.0.1]
+      --port <PORT>                [default: 3001]
+      --worker-urls <WORKER_URLS>  
+      --policy <POLICY>            [default: round_robin] [possible values: round_robin, random]
+  -h, --help                       Print help
+  -V, --version                    Print version
+```
+
+### Setting Up Workers
+
+1. Launch worker instances
+```bash
+# Launch first worker on GPU 0
+export CUDA_VISIBLE_DEVICES=0
+python -m sglang.launch_server \
+    --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \
+    --host 127.0.0.1 \
+    --port 30000
+
+# Launch second worker on GPU 1
+export CUDA_VISIBLE_DEVICES=1
+python -m sglang.launch_server \
+    --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \
+    --host 127.0.0.1 \
+    --port 30002
+```
+
+2. Launch router and connect to workers
+```bash
+./target/debug/router --worker-urls http://127.0.0.1:30000,http://127.0.0.1:30002
+```
+
+**Note**: This module is still experimental. Please expect active changes and updates.
+
+### Python bindings
+
+```bash
+$ cargo build --release
+$ maturin build -i /usr/bin/python
+$ pip install <path to wheel>
+```

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),
+    }
+}