49 lines
1012 B
Rust
49 lines
1012 B
Rust
use pyo3::prelude::*;
|
|
pub mod router;
|
|
mod server;
|
|
pub mod tree;
|
|
|
|
// 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(())
|
|
}
|