lib.rs 2.5 KB
Newer Older
1
2
use pyo3::prelude::*;
pub mod router;
3
pub mod server;
4
pub mod tree;
5

6
7
8
9
10
#[pyclass(eq)]
#[derive(Clone, PartialEq)]
pub enum PolicyType {
    Random,
    RoundRobin,
11
    CacheAware,
12
13
}

14
15
16
17
18
#[pyclass]
struct Router {
    host: String,
    port: u16,
    worker_urls: Vec<String>,
19
    policy: PolicyType,
20
    cache_threshold: f32,
21
22
    balance_abs_threshold: usize,
    balance_rel_threshold: f32,
23
24
    eviction_interval_secs: u64,
    max_tree_size: usize,
25
26
27
28
29
}

#[pymethods]
impl Router {
    #[new]
30
31
32
33
34
    #[pyo3(signature = (
        worker_urls,
        policy = PolicyType::RoundRobin,
        host = String::from("127.0.0.1"),
        port = 3001,
35
        cache_threshold = 0.50,
36
37
        balance_abs_threshold = 32,
        balance_rel_threshold = 1.0001,
38
39
        eviction_interval_secs = 60,
        max_tree_size = 2usize.pow(24)
40
41
42
43
44
45
    ))]
    fn new(
        worker_urls: Vec<String>,
        policy: PolicyType,
        host: String,
        port: u16,
46
        cache_threshold: f32,
47
48
        balance_abs_threshold: usize,
        balance_rel_threshold: f32,
49
50
        eviction_interval_secs: u64,
        max_tree_size: usize,
51
52
    ) -> PyResult<Self> {
        Ok(Router {
53
54
55
            host,
            port,
            worker_urls,
56
            policy,
57
            cache_threshold,
58
59
            balance_abs_threshold,
            balance_rel_threshold,
60
61
            eviction_interval_secs,
            max_tree_size,
62
        })
63
64
65
66
67
68
    }

    fn start(&self) -> PyResult<()> {
        let host = self.host.clone();
        let port = self.port;
        let worker_urls = self.worker_urls.clone();
69
70
71
72

        let policy_config = match &self.policy {
            PolicyType::Random => router::PolicyConfig::RandomConfig,
            PolicyType::RoundRobin => router::PolicyConfig::RoundRobinConfig,
73
74
            PolicyType::CacheAware => router::PolicyConfig::CacheAwareConfig {
                cache_threshold: self.cache_threshold,
75
76
                balance_abs_threshold: self.balance_abs_threshold,
                balance_rel_threshold: self.balance_rel_threshold,
77
78
                eviction_interval_secs: self.eviction_interval_secs,
                max_tree_size: self.max_tree_size,
79
80
            },
        };
81
82

        actix_web::rt::System::new().block_on(async move {
83
            server::startup(host, port, worker_urls, policy_config)
84
85
                .await
                .unwrap();
86
87
88
89
90
91
92
        });

        Ok(())
    }
}

#[pymodule]
93
fn sglang_router_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
94
    m.add_class::<PolicyType>()?;
95
96
    m.add_class::<Router>()?;
    Ok(())
97
}