[router] Expose worker startup secs & Return error instead of panic for router init (#3016)

This commit is contained in:
Byron Hsu
2025-01-20 12:45:13 -08:00
committed by GitHub
parent 5dfcacfcb1
commit 0311ce8e1c
7 changed files with 124 additions and 47 deletions

View File

@@ -17,6 +17,7 @@ struct Router {
port: u16,
worker_urls: Vec<String>,
policy: PolicyType,
worker_startup_timeout_secs: u64,
cache_threshold: f32,
balance_abs_threshold: usize,
balance_rel_threshold: f32,
@@ -34,6 +35,7 @@ impl Router {
policy = PolicyType::RoundRobin,
host = String::from("127.0.0.1"),
port = 3001,
worker_startup_timeout_secs = 300,
cache_threshold = 0.50,
balance_abs_threshold = 32,
balance_rel_threshold = 1.0001,
@@ -47,6 +49,7 @@ impl Router {
policy: PolicyType,
host: String,
port: u16,
worker_startup_timeout_secs: u64,
cache_threshold: f32,
balance_abs_threshold: usize,
balance_rel_threshold: f32,
@@ -60,6 +63,7 @@ impl Router {
port,
worker_urls,
policy,
worker_startup_timeout_secs,
cache_threshold,
balance_abs_threshold,
balance_rel_threshold,
@@ -72,9 +76,14 @@ impl Router {
fn start(&self) -> PyResult<()> {
let policy_config = match &self.policy {
PolicyType::Random => router::PolicyConfig::RandomConfig,
PolicyType::RoundRobin => router::PolicyConfig::RoundRobinConfig,
PolicyType::Random => router::PolicyConfig::RandomConfig {
timeout_secs: self.worker_startup_timeout_secs,
},
PolicyType::RoundRobin => router::PolicyConfig::RoundRobinConfig {
timeout_secs: self.worker_startup_timeout_secs,
},
PolicyType::CacheAware => router::PolicyConfig::CacheAwareConfig {
timeout_secs: self.worker_startup_timeout_secs,
cache_threshold: self.cache_threshold,
balance_abs_threshold: self.balance_abs_threshold,
balance_rel_threshold: self.balance_rel_threshold,
@@ -93,10 +102,9 @@ impl Router {
max_payload_size: self.max_payload_size,
})
.await
.unwrap();
});
Ok(())
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
Ok(())
})
}
}