diff --git a/crates/aisix-admin/src/openapi.rs b/crates/aisix-admin/src/openapi.rs index b7334329..8fb922ad 100644 --- a/crates/aisix-admin/src/openapi.rs +++ b/crates/aisix-admin/src/openapi.rs @@ -3730,7 +3730,13 @@ fn add_variant_titles(doc: &mut Value) { ), ( "/components/schemas/RoutingStrategy/oneOf", - &["Round robin", "Weighted", "Failover", "Least cost"], + &[ + "Round robin", + "Weighted", + "Failover", + "Least cost", + "Least latency", + ], ), ( "/components/schemas/SlsContentMode/oneOf", diff --git a/crates/aisix-core/src/models/routing.rs b/crates/aisix-core/src/models/routing.rs index 2eaa8ad9..7c41cc75 100644 --- a/crates/aisix-core/src/models/routing.rs +++ b/crates/aisix-core/src/models/routing.rs @@ -18,7 +18,12 @@ //! attempt them best-first, falling forward down the ranked order: //! - `least_cost`: cheapest target first, by the target model's `cost` //! (combined input+output per-1K price). Targets without a `cost` rank -//! last. See [`RoutingStrategy::is_metric_based`]. +//! last. +//! - `least_latency`: fastest target first, by a moving average of recent +//! observed upstream latency (time-to-first-token for streaming). Targets +//! with no latency samples yet rank first so they get probed. +//! +//! See [`RoutingStrategy::is_metric_based`]. use serde::{Deserialize, Serialize}; @@ -39,6 +44,10 @@ pub enum RoutingStrategy { /// input+output per-1K price), then fall forward. Targets without a /// configured `cost` rank last. LeastCost, + /// Rank targets fastest-first by a moving average of recent observed + /// upstream latency (time-to-first-token for streaming), then fall + /// forward. Targets with no samples yet rank first so they get probed. + LeastLatency, } impl RoutingStrategy { @@ -47,7 +56,10 @@ impl RoutingStrategy { /// strategies are ordered after target resolution, where each target's /// Model and runtime state are available. pub fn is_metric_based(&self) -> bool { - matches!(self, RoutingStrategy::LeastCost) + matches!( + self, + RoutingStrategy::LeastCost | RoutingStrategy::LeastLatency + ) } } @@ -243,17 +255,23 @@ mod tests { } #[test] - fn parses_least_cost_strategy() { - let r: Routing = serde_json::from_str( + fn parses_metric_strategies() { + let cost: Routing = serde_json::from_str( r#"{"strategy":"least_cost","targets":[{"model":"a"},{"model":"b"}]}"#, ) .unwrap(); - assert_eq!(r.strategy, RoutingStrategy::LeastCost); + assert_eq!(cost.strategy, RoutingStrategy::LeastCost); + let latency: Routing = serde_json::from_str( + r#"{"strategy":"least_latency","targets":[{"model":"a"},{"model":"b"}]}"#, + ) + .unwrap(); + assert_eq!(latency.strategy, RoutingStrategy::LeastLatency); } #[test] fn is_metric_based_classification() { assert!(RoutingStrategy::LeastCost.is_metric_based()); + assert!(RoutingStrategy::LeastLatency.is_metric_based()); assert!(!RoutingStrategy::Failover.is_metric_based()); assert!(!RoutingStrategy::RoundRobin.is_metric_based()); assert!(!RoutingStrategy::Weighted.is_metric_based()); diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index a3e94824..7bd56b37 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1080,6 +1080,10 @@ async fn dispatch( Ok(upstream) => { state.health.record_success(&model.display_name); state.runtime_status.mark_healthy(&attempt.id); + // Feed the least_latency EWMA. For streaming this is + // time-to-first-response (upstream stream established) — + // the routing-relevant latency signal. + state.runtime_status.record_latency(&attempt.id, latency_ms); stream_routing.attempts.push(AttemptRecord { index: idx, kind, @@ -1706,6 +1710,11 @@ async fn dispatch( Ok(resp) => { state.health.record_success(&model.display_name); state.runtime_status.mark_healthy(&attempt.id); + // Feed the least_latency EWMA with this attempt's + // round-trip latency. + state + .runtime_status + .record_latency(&attempt.id, attempt_latency_ms); chosen_provider = Some(provider.to_ascii_lowercase()); chosen_provider_key_id = Some(pk_entry.id.clone()); chosen_upstream_model = diff --git a/crates/aisix-proxy/src/health.rs b/crates/aisix-proxy/src/health.rs index 49c2323e..b0ce94ff 100644 --- a/crates/aisix-proxy/src/health.rs +++ b/crates/aisix-proxy/src/health.rs @@ -225,6 +225,10 @@ struct RuntimeEntry { last_checked_at: Option, last_check_status: Option, status_reason: Option, + /// Exponentially-weighted moving average of recent observed upstream + /// latency in milliseconds. `None` until the first sample. Drives the + /// `least_latency` routing strategy; independent of health/cooldown. + latency_ewma_ms: Option, } impl RuntimeEntry { @@ -332,6 +336,12 @@ pub struct HealthTracker { entries: DashMap, } +/// Smoothing factor for the per-target latency EWMA. Higher = more weight on +/// the most recent sample (faster reaction to a slowing upstream), lower = +/// smoother. 0.3 balances reacting to a real regression against per-request +/// jitter, roughly matching LiteLLM's last-10-samples moving average. +const LATENCY_EWMA_ALPHA: f64 = 0.3; + #[derive(Default, Debug)] pub struct ModelRuntimeStatusTracker { entries: DashMap, @@ -455,6 +465,30 @@ impl ModelRuntimeStatusTracker { }); } + /// Fold a fresh latency sample (ms) into the target's EWMA. Called on + /// each successful upstream attempt; drives the `least_latency` routing + /// strategy. Independent of health/cooldown state. + pub fn record_latency(&self, model_id: &str, latency_ms: u32) { + let sample = f64::from(latency_ms); + self.entries + .entry(model_id.to_string()) + .and_modify(|entry| { + entry.latency_ewma_ms = Some(match entry.latency_ewma_ms { + Some(prev) => LATENCY_EWMA_ALPHA * sample + (1.0 - LATENCY_EWMA_ALPHA) * prev, + None => sample, + }); + }) + .or_insert_with(|| RuntimeEntry { + latency_ewma_ms: Some(sample), + ..RuntimeEntry::default() + }); + } + + /// Current latency EWMA (ms) for `model_id`, or `None` if never sampled. + pub fn latency_ewma_ms(&self, model_id: &str) -> Option { + self.entries.get(model_id).and_then(|e| e.latency_ewma_ms) + } + pub fn status(&self, model_id: &str) -> RuntimeStatusSnapshot { self.status_with_stale(model_id, None) } diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 700912d1..94bbf3a2 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -566,6 +566,9 @@ async fn dispatch( .await { Ok(mut outcome) => { + let latency_ms = ms_since(attempt_started); + // Feed the least_latency EWMA for this target. + state.runtime_status.record_latency(&target.id, latency_ms); routing.attempts.push(AttemptRecord { index: idx, kind, @@ -576,7 +579,7 @@ async fn dispatch( success: true, error_class: String::new(), error_message: String::new(), - latency_ms: ms_since(attempt_started), + latency_ms, }); outcome.routing = routing; return Ok(outcome); diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 07e66cf9..745ede57 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -446,6 +446,9 @@ async fn dispatch( }; match result { Ok(mut success) => { + let latency_ms = ms_since(attempt_started); + // Feed the least_latency EWMA for this target. + state.runtime_status.record_latency(&target.id, latency_ms); routing.attempts.push(AttemptRecord { index: idx, kind, @@ -456,7 +459,7 @@ async fn dispatch( success: true, error_class: String::new(), error_message: String::new(), - latency_ms: ms_since(attempt_started), + latency_ms, }); success.routing = routing; return Ok(success); diff --git a/crates/aisix-proxy/src/routing.rs b/crates/aisix-proxy/src/routing.rs index 1a957270..b406f048 100644 --- a/crates/aisix-proxy/src/routing.rs +++ b/crates/aisix-proxy/src/routing.rs @@ -138,7 +138,7 @@ impl RoutingRegistry { RoutingStrategy::Weighted => weighted_pick(&routing.targets), // Metric-ordered strategies never reach here — `pick_targets` // short-circuits them before computing a start index. - RoutingStrategy::LeastCost => 0, + RoutingStrategy::LeastCost | RoutingStrategy::LeastLatency => 0, } } @@ -216,15 +216,33 @@ fn cost_key(model: &Model) -> f64 { .unwrap_or(f64::INFINITY) } +/// Observed-latency key used to rank `least_latency` targets. A target with +/// no latency samples yet sorts first (treated as −∞) so it gets probed; +/// once it has an EWMA it ranks by that. +fn latency_key(runtime_status: &crate::ModelRuntimeStatusTracker, id: &str) -> f64 { + runtime_status + .latency_ewma_ms(id) + .unwrap_or(f64::NEG_INFINITY) +} + /// Rank the resolved attempt list by the strategy's runtime metric, /// best-first (ascending). Stable, so equal-metric targets keep their /// declaration order. Only metric-based strategies reach here; positional /// strategies are ordered in [`RoutingRegistry::pick_targets`]. -fn order_attempts_by_metric(strategy: RoutingStrategy, attempts: &mut [AttemptModel]) { +fn order_attempts_by_metric( + strategy: RoutingStrategy, + attempts: &mut [AttemptModel], + runtime_status: &crate::ModelRuntimeStatusTracker, +) { match strategy { RoutingStrategy::LeastCost => { attempts.sort_by(|a, b| cost_key(&a.model).total_cmp(&cost_key(&b.model))); } + RoutingStrategy::LeastLatency => { + attempts.sort_by(|a, b| { + latency_key(runtime_status, &a.id).total_cmp(&latency_key(runtime_status, &b.id)) + }); + } RoutingStrategy::Failover | RoutingStrategy::RoundRobin | RoutingStrategy::Weighted => {} } } @@ -359,7 +377,7 @@ pub(crate) fn resolve_attempt_models( // rank it best-first here (target Models are now resolved) and cap it to // the same attempt budget the positional strategies apply upstream. if routing.strategy.is_metric_based() { - order_attempts_by_metric(routing.strategy, &mut resolved); + order_attempts_by_metric(routing.strategy, &mut resolved, runtime_status); resolved.truncate(routing.max_fallbacks_or_default() + 1); } match filter_attempt_models( @@ -756,24 +774,26 @@ mod tests { #[test] fn least_cost_orders_cheapest_first() { + let t = crate::ModelRuntimeStatusTracker::new(); let mut attempts = vec![ am_with_cost("pricey", 10.0, 20.0), // 30 / 1K am_with_cost("cheap", 1.0, 2.0), // 3 / 1K am_with_cost("mid", 5.0, 5.0), // 10 / 1K ]; - order_attempts_by_metric(RoutingStrategy::LeastCost, &mut attempts); + order_attempts_by_metric(RoutingStrategy::LeastCost, &mut attempts, &t); let ids: Vec<&str> = attempts.iter().map(|a| a.id.as_str()).collect(); assert_eq!(ids, vec!["cheap", "mid", "pricey"]); } #[test] fn least_cost_ranks_missing_cost_last_and_stably() { + let t = crate::ModelRuntimeStatusTracker::new(); let mut attempts = vec![ am("no-cost-a"), // +∞ am_with_cost("cheap", 1.0, 1.0), // 2 / 1K am("no-cost-b"), // +∞ ]; - order_attempts_by_metric(RoutingStrategy::LeastCost, &mut attempts); + order_attempts_by_metric(RoutingStrategy::LeastCost, &mut attempts, &t); let ids: Vec<&str> = attempts.iter().map(|a| a.id.as_str()).collect(); // Priced target first; equal (missing-cost) targets keep their // declaration order thanks to the stable sort. @@ -782,12 +802,49 @@ mod tests { #[test] fn non_metric_strategy_leaves_order_untouched() { + let t = crate::ModelRuntimeStatusTracker::new(); let mut attempts = vec![am_with_cost("b", 9.0, 9.0), am_with_cost("a", 1.0, 1.0)]; - order_attempts_by_metric(RoutingStrategy::Failover, &mut attempts); + order_attempts_by_metric(RoutingStrategy::Failover, &mut attempts, &t); let ids: Vec<&str> = attempts.iter().map(|a| a.id.as_str()).collect(); assert_eq!(ids, vec!["b", "a"]); } + // ── order_attempts_by_metric (least_latency) ────────────────── + #[test] + fn least_latency_orders_fastest_first() { + let t = crate::ModelRuntimeStatusTracker::new(); + t.record_latency("slow", 900); + t.record_latency("fast", 50); + t.record_latency("mid", 300); + let mut attempts = vec![am("slow"), am("fast"), am("mid")]; + order_attempts_by_metric(RoutingStrategy::LeastLatency, &mut attempts, &t); + let ids: Vec<&str> = attempts.iter().map(|a| a.id.as_str()).collect(); + assert_eq!(ids, vec!["fast", "mid", "slow"]); + } + + #[test] + fn least_latency_probes_unmeasured_targets_first() { + let t = crate::ModelRuntimeStatusTracker::new(); + t.record_latency("measured", 100); + // "unseen-a"/"unseen-b" have no samples → rank first (−∞), keeping + // their declaration order via the stable sort. + let mut attempts = vec![am("measured"), am("unseen-a"), am("unseen-b")]; + order_attempts_by_metric(RoutingStrategy::LeastLatency, &mut attempts, &t); + let ids: Vec<&str> = attempts.iter().map(|a| a.id.as_str()).collect(); + assert_eq!(ids, vec!["unseen-a", "unseen-b", "measured"]); + } + + #[test] + fn record_latency_ewma_tracks_recent_samples() { + let t = crate::ModelRuntimeStatusTracker::new(); + assert_eq!(t.latency_ewma_ms("m"), None); + t.record_latency("m", 100); + assert_eq!(t.latency_ewma_ms("m"), Some(100.0)); // first sample seeds + t.record_latency("m", 200); + // 0.3*200 + 0.7*100 = 130 + assert!((t.latency_ewma_ms("m").unwrap() - 130.0).abs() < 1e-9); + } + #[test] fn metric_strategy_pick_targets_returns_full_declaration_order() { let reg = RoutingRegistry::new(); diff --git a/schemas/resources/model.schema.json b/schemas/resources/model.schema.json index a1e90e23..591223d8 100644 --- a/schemas/resources/model.schema.json +++ b/schemas/resources/model.schema.json @@ -436,6 +436,13 @@ "least_cost" ], "type": "string" + }, + { + "description": "Rank targets fastest-first by a moving average of recent observed upstream latency (time-to-first-token for streaming), then fall forward. Targets with no samples yet rank first so they get probed.", + "enum": [ + "least_latency" + ], + "type": "string" } ] }, diff --git a/schemas/resources/routing.schema.json b/schemas/resources/routing.schema.json index f6c2c6a7..566f6b0e 100644 --- a/schemas/resources/routing.schema.json +++ b/schemas/resources/routing.schema.json @@ -91,6 +91,13 @@ "enum": [ "least_cost" ] + }, + { + "description": "Rank targets fastest-first by a moving average of recent observed upstream latency (time-to-first-token for streaming), then fall forward. Targets with no samples yet rank first so they get probed.", + "type": "string", + "enum": [ + "least_latency" + ] } ] }, diff --git a/tests/e2e/src/cases/latency-aware-routing-e2e.test.ts b/tests/e2e/src/cases/latency-aware-routing-e2e.test.ts new file mode 100644 index 00000000..cc5540b3 --- /dev/null +++ b/tests/e2e/src/cases/latency-aware-routing-e2e.test.ts @@ -0,0 +1,161 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +const CALLER_PLAINTEXT = "sk-latency-routing-e2e-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +function okBody(content: string) { + return { + id: `cmpl-${content}`, + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }; +} + +describe("latency-aware (least_latency) routing e2e", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + const upstreams: OpenAiUpstream[] = []; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["*"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await Promise.all(upstreams.map((u) => u.close())); + }); + + async function createOpenAiModel( + displayName: string, + upstream: OpenAiUpstream, + extra: Record = {}, + ): Promise { + if (!admin) throw new Error("admin client not initialized"); + const providerKey = await admin.createProviderKey({ + display_name: `${displayName}-pk`, + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: displayName, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: providerKey.id, + ...extra, + }); + } + + function client(): OpenAI { + return new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app?.proxyUrl}/v1`, + maxRetries: 0, + }); + } + + test("routes to the fastest target once latencies are learned", async (ctx) => { + if (!etcdReachable || !app || !admin) { + ctx.skip(); + return; + } + + const slow = await startOpenAiUpstream({ + responseDelayMs: 400, + nonStreamBody: okBody("slow-served"), + }); + const fast = await startOpenAiUpstream({ nonStreamBody: okBody("fast-served") }); + upstreams.push(slow, fast); + + await createOpenAiModel("lat-slow", slow); + await createOpenAiModel("lat-fast", fast); + // Declare the slow target FIRST — least_latency must reorder by observed + // latency, not honor declaration order. + await admin.createModel({ + display_name: "lat-virtual", + routing: { + strategy: "least_latency", + targets: [{ model: "lat-slow" }, { model: "lat-fast" }], + }, + }); + + const c = client(); + // Wait for the virtual to dispatch through the DP snapshot. The probe is + // expected to 404 until the watch loop applies the record, so a failed + // attempt legitimately retries here. This also seeds cold-start latency + // samples (unmeasured targets are tried in declaration order first). + await waitConfigPropagation(async () => { + try { + const p = await c.chat.completions.create({ + model: "lat-virtual", + messages: [{ role: "user", content: "warmup" }], + }); + return ["slow-served", "fast-served"].includes( + p.choices[0]?.message.content ?? "", + ); + } catch { + return false; + } + }); + // Extra warmup so BOTH targets definitely have a latency EWMA (cold-start + // tries them in declaration order, one per request). + for (let i = 0; i < 3; i++) { + await c.chat.completions.create({ + model: "lat-virtual", + messages: [{ role: "user", content: `warm-${i}` }], + }); + } + + const slowBaseline = slow.receivedRequests.length; + const fastBaseline = fast.receivedRequests.length; + const contents: string[] = []; + for (let i = 0; i < 5; i++) { + const r = await c.chat.completions.create({ + model: "lat-virtual", + messages: [{ role: "user", content: `measure-${i}` }], + }); + contents.push(r.choices[0]?.message.content ?? ""); + } + + // Steady state: every request now goes to the fast target. + expect(contents).toEqual([ + "fast-served", + "fast-served", + "fast-served", + "fast-served", + "fast-served", + ]); + expect(fast.receivedRequests.length - fastBaseline).toBe(5); + expect(slow.receivedRequests.length - slowBaseline).toBe(0); + }); +});