diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index 73b049abf..ef721eb3c 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/sgl-model-gateway/src/routers/http/router.rs @@ -205,8 +205,7 @@ impl Router { }; let load_incremented = if policy.name() == "cache_aware" { - worker.increment_load(); - RouterMetrics::set_running_requests(worker.url(), worker.load()); + increment_load(&worker); true } else { false @@ -246,11 +245,7 @@ impl Router { // won't have done it (it only decrements on success or non-retryable failures) if is_retryable_status(response.status()) && load_incremented { if let Some(cleanup_worker) = worker_for_cleanup { - cleanup_worker.decrement_load(); - RouterMetrics::set_running_requests( - cleanup_worker.url(), - cleanup_worker.load(), - ); + decrement_load(&cleanup_worker); } } @@ -518,8 +513,7 @@ impl Router { // Decrement load on error if it was incremented if load_incremented { if let Some(ref w) = worker { - w.decrement_load(); - RouterMetrics::set_running_requests(worker_url, w.load()); + decrement_load(w); } } @@ -551,8 +545,7 @@ impl Router { // IMPORTANT: Decrement load on error before returning if load_incremented { if let Some(ref w) = worker { - w.decrement_load(); - RouterMetrics::set_running_requests(worker_url, w.load()); + decrement_load(w); } } @@ -564,8 +557,7 @@ impl Router { // Decrement load counter for non-streaming requests if it was incremented if load_incremented { if let Some(ref w) = worker { - w.decrement_load(); - RouterMetrics::set_running_requests(worker_url, w.load()); + decrement_load(w); } } @@ -594,11 +586,7 @@ impl Router { // Check for stream end marker using memmem for efficiency if memmem::find(&bytes, b"data: [DONE]").is_some() { if let Some(ref w) = stream_worker { - w.decrement_load(); - RouterMetrics::set_running_requests( - &worker_url_owned, - w.load(), - ); + decrement_load(w); decremented = true; } } @@ -614,8 +602,7 @@ impl Router { } if !decremented { if let Some(ref w) = stream_worker { - w.decrement_load(); - RouterMetrics::set_running_requests(&worker_url_owned, w.load()); + decrement_load(w); } } }); @@ -685,6 +672,16 @@ impl Router { } } +fn increment_load(w: &Arc) { + w.increment_load(); + RouterMetrics::set_running_requests(w.url(), w.load()); +} + +fn decrement_load(w: &Arc) { + w.decrement_load(); + RouterMetrics::set_running_requests(w.url(), w.load()); +} + use async_trait::async_trait; #[async_trait]