diff --git a/litellm-rust/crates/router-core/src/call.rs b/litellm-rust/crates/router-core/src/call.rs index 3a30c779958..a118cac28c1 100644 --- a/litellm-rust/crates/router-core/src/call.rs +++ b/litellm-rust/crates/router-core/src/call.rs @@ -4,13 +4,7 @@ //! `should_retry_this_error` and `_time_to_sleep_before_retry` in `litellm/router.py`, and //! `run_async_fallback` in `router_utils/fallback_event_handlers.py`. -use std::{ - borrow::Cow, - future::Future, - pin::Pin, - sync::Arc, - time::{Duration, Instant}, -}; +use std::{borrow::Cow, future::Future, pin::Pin, sync::Arc, time::Duration}; use serde_json::{Map, Value}; @@ -1231,7 +1225,7 @@ impl Core { .unrouted_attempt(snapshot, op, request, target, (retry, max_retries), at) .await; } - let pick_started = Instant::now(); + let pick_started = self.health.now_time(); let (deployment, candidates) = match self .picked(snapshot, op, request, target, inputs, retry, at) .await @@ -1248,7 +1242,11 @@ impl Core { self.host.observe(&RouterEvent::Selected { model_group: candidates.model_group.clone(), deployment_id: deployment.id.clone(), - pick: pick_started.elapsed(), + pick: self + .health + .now_time() + .duration_since(pick_started) + .unwrap_or_default(), }); if !self.host.supports(op, deployment.provider()) { return Err(Error::Config(ConfigError::Unsupported(Unsupported::Op { @@ -1274,7 +1272,7 @@ impl Core { }, None, snapshot.settings.dynamic.default_ref(), - jiff::Timestamp::now(), + self.health.timestamp(), ), Ok(None) ); diff --git a/litellm-rust/crates/router-core/src/health.rs b/litellm-rust/crates/router-core/src/health.rs index c8055eafe31..eaa61dac7ba 100644 --- a/litellm-rust/crates/router-core/src/health.rs +++ b/litellm-rust/crates/router-core/src/health.rs @@ -247,6 +247,10 @@ impl Health { self.clock.now() } + pub(crate) fn timestamp(&self) -> jiff::Timestamp { + jiff::Timestamp::try_from(self.clock.now()).unwrap_or(jiff::Timestamp::UNIX_EPOCH) + } + fn now(&self) -> f64 { unix_seconds(self.clock.now()) } diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs index b6ada42df1f..c7f3b1f9c31 100644 --- a/litellm-rust/crates/router-core/src/select.rs +++ b/litellm-rust/crates/router-core/src/select.rs @@ -312,7 +312,7 @@ impl Core { caller, default_target, snapshot.settings.dynamic.default_ref(), - jiff::Timestamp::now(), + self.health.timestamp(), )? .ok_or_else(|| Error::UnknownModel { // The name `_common_checks_available_deployment` looked for, past the alias. diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs index 9f41371f255..a2a741d6135 100644 --- a/litellm-rust/crates/router-core/tests/call_loop.rs +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -1755,3 +1755,57 @@ async fn a_listed_pass_through_pick_names_the_group_an_alias_resolves_to(config: Error::NoPassThroughDeployments { model } if model == "primary" )); } + +struct StoppedClock(u64); + +impl litellm_router_core::health::Clock for StoppedClock { + fn now(&self) -> std::time::SystemTime { + std::time::UNIX_EPOCH + Duration::from_secs(self.0) + } +} + +fn utc_hour(hour: u64) -> u64 { + 1_767_225_600 + hour * 3_600 +} + +#[rstest] +#[case::inside_the_window(utc_hour(10), vec!["open"])] +#[case::outside_the_window(utc_hour(20), vec!["open", "reserved"])] +#[tokio::test] +async fn access_windows_follow_the_routers_clock(#[case] now: u64, #[case] expected: Vec<&str>) { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "chat", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "open"}}, + {"model_name": "chat", "litellm_params": {"model": "openai/b"}, "model_info": { + "id": "reserved", + "access_windows": [{"start": "09:00", "end": "17:00", "timezone": "UTC", "team_ids": ["batch"]}], + }}, + ] + })) + .unwrap(); + let router = Router::builder( + config, + RouterDynamicConfig::default(), + Arc::new(ScriptedHost::default()), + ) + .randomness(Arc::new(SeededRandomness::new(7))) + .clock(Arc::new(StoppedClock(now))) + .build() + .unwrap(); + let mut picked = std::collections::BTreeSet::new(); + for _ in 0..40 { + let deployment = router + .select(&Op::from("acompletion"), &request_for("chat")) + .await + .unwrap(); + picked.insert(deployment.id.clone()); + } + assert_eq!(picked.into_iter().collect::>(), expected); +} + +fn request_for(model: &str) -> RouterRequest { + RouterRequest { + model: model.to_owned(), + ..RouterRequest::default() + } +}