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

This commit is contained in:
Byron Hsu
2024-10-28 09:49:48 -07:00
committed by GitHub
parent 6e13b650a9
commit 3839be2913
8 changed files with 2666 additions and 0 deletions

45
rust/src/lib.rs Normal file
View File

@@ -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(())
}