From 77a568da4446310a4634681b36ddc841363dfb9c Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 13:51:09 -0700 Subject: [PATCH 1/6] feat(stargate): route returning sessions to their last cluster Add last_cluster_affinity to wait-and-widen and pulsar-wait-and-widen. When enabled, Stargate classifies each request with an affinity key once, on its first load-balancer decision, from the trusted x-stargate-last-cluster-id header: - new (no hint) and stale (hint names a cluster outside the candidate set) requests skip the affinity wait and the affinity prefill discount and go to the best-ranked cluster that is selectable now. - returning requests move the last cluster to the front of the affinity order for that request only, then run the configured wait, discount, and widening unchanged. Per-key ring and ranking caches keep the original order. Other algorithms reject the field at startup through the existing unsupported-field error. Routing expressions do not accept it. The header is consumed and never forwarded to Pylon. Remove the pulsar-wait-and-widen eligible-primary shortcut that sent a request to its rank-1 cluster without a load check when no queue-SLO fields were set. Rank 1 now goes through queue admission, the affinity wait, and the discount like any other request. Add stargate_routing_session_selections_total, the routing.session_state span field, and a debug log for each classified selection. Add unit and integration tests, including the cache-thrash regression for both algorithms with the flag on and off, a pulsar-wait-and-widen load-balancer microbenchmark case, and docs. Refs #2289 Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../metrics/llm-request-router/metrics.md | 7 +- .../crates/protocol/src/tunnel_contract.rs | 9 + .../stargate-bench/src/microbench/lb.rs | 84 +- .../crates/stargate-routing-sim/src/sim.rs | 1 + .../crates/stargate/src/http_proxy/attempt.rs | 2 + .../crates/stargate/src/http_proxy/request.rs | 56 +- .../crates/stargate/src/http_proxy/routing.rs | 1 + .../crates/stargate/src/http_proxy/run.rs | 13 + .../crates/stargate/src/http_proxy/trace.rs | 85 ++ .../stargate/src/http_proxy/upstream.rs | 1 + .../stargate/src/load_balancer/config.rs | 5 + .../stargate/src/load_balancer/expression.rs | 8 + .../crates/stargate/src/load_balancer/mod.rs | 3 + .../stargate/src/load_balancer/power_of_n.rs | 3 + .../load_balancer/pulsar_wait_and_widen.rs | 704 ++++++++++++++-- .../stargate/src/load_balancer/random.rs | 1 + .../stargate/src/load_balancer/request.rs | 3 + .../stargate/src/load_balancer/round_robin.rs | 1 + .../stargate/src/load_balancer/session.rs | 759 ++++++++++++++++++ .../stargate/src/load_balancer/tests.rs | 693 +++++++++++++++- .../src/load_balancer/wait_and_widen.rs | 31 +- .../stargate/crates/stargate/src/metrics.rs | 2 + .../stargate/src/routing_state/tests.rs | 4 + .../stargate/tests/suite/load_balancing.rs | 276 +++++++ .../stargate/tests/suite/proxy_contract.rs | 180 +++++ .../docs/load-balancer-configuration.md | 179 ++++- 26 files changed, 3016 insertions(+), 95 deletions(-) create mode 100644 src/libraries/rust/stargate/crates/stargate/src/load_balancer/session.rs diff --git a/docs/observability/metrics/llm-request-router/metrics.md b/docs/observability/metrics/llm-request-router/metrics.md index 29c85e2532..c6e74c5c07 100644 --- a/docs/observability/metrics/llm-request-router/metrics.md +++ b/docs/observability/metrics/llm-request-router/metrics.md @@ -14,9 +14,9 @@ The self-managed stack maps ## Label Boundaries Use bounded labels only. Keep `routing_key`, `model`, `inference_server_id`, -`algorithm`, `selection`, `status`, `result`, `reason`, and `cache` to bounded -service dimensions. Routing expression content never becomes a -label value. Do not add request IDs, session IDs, function IDs, organization +`algorithm`, `selection`, `session_state`, `status`, `result`, `reason`, and +`cache` to bounded service dimensions. Routing expression content never becomes +a label value. Do not add request IDs, session IDs, function IDs, organization IDs, project IDs, raw URLs, raw prompts, authorization values, or other unbounded request fields as metric labels. @@ -29,6 +29,7 @@ unbounded request fields as metric labels. | `stargate_proxy_retries_total` | Counter | `llm-request-router:9090/metrics` | `routing_key`, `model`, `reason` | Total proxy retries by retry reason. | | `stargate_routing_selections_total` | Counter | `llm-request-router:9090/metrics` | `routing_key`, `model`, `algorithm`, `selection` | Primary and ranked fallback cluster choices used for upstream attempts. | | `stargate_routing_kv_free_token_fallback_selections_total` | Counter | `llm-request-router:9090/metrics` | `routing_key`, `model`, `algorithm` | Routes selected after a higher-ranked candidate failed the KV free-token check. | +| `stargate_routing_session_selections_total` | Counter | `llm-request-router:9090/metrics` | `routing_key`, `model`, `algorithm`, `session_state`, `selection` | Cluster choices for requests classified by `last_cluster_affinity`. `session_state` is `new`, `returning`, or `stale`. For `returning`, `primary` means the request went to its last cluster; for `new` and `stale`, it means the cluster came from the affinity group. Every state and selection pair starts at zero once a routing key, model, and algorithm records its first classified request. | | `stargate_routing_expressions_total` | Counter | `llm-request-router:9090/metrics` | `algorithm`, `cache` | Accepted `x-routing-method` values with parameters. `cache` is `hit`, `build`, or `rebuild`. Values without parameters and rejected values are not counted. Rejected `x-routing-method` values write the warn log `invalid routing algorithm header` with their `rejection_reason`. | | `stargate_routing_expression_cache_entries` | Gauge | `llm-request-router:9090/metrics` | none | Approximate number of routing targets that hold a routing expression configuration, with a capacity of 16384 per process. The router runs cache maintenance once a minute to remove entries idle for 15 minutes. Scrapes read the count without maintenance, so it can lag behind pending insertions or removals until the next maintenance run. | | `stargate_proxy_retry_exhausted_total` | Counter | `llm-request-router:9090/metrics` | `routing_key`, `model`, `reason` | Total requests that exhausted retry options. | diff --git a/src/libraries/rust/stargate/crates/protocol/src/tunnel_contract.rs b/src/libraries/rust/stargate/crates/protocol/src/tunnel_contract.rs index 0c2bde3c05..b06592c97c 100644 --- a/src/libraries/rust/stargate/crates/protocol/src/tunnel_contract.rs +++ b/src/libraries/rust/stargate/crates/protocol/src/tunnel_contract.rs @@ -24,6 +24,8 @@ pub const HEADER_STARGATE_UPSTREAM_RETRYABLE: &str = "x-stargate-upstream-retrya pub const HEADER_STARGATE_RETRYABLE: &str = "x-stargate-retryable"; pub const HEADER_STARGATE_RETRY_REASON: &str = "x-stargate-retry-reason"; pub const HEADER_STARGATE_RETRY_AFTER_MS: &str = "x-stargate-retry-after-ms"; +/// Gateway hint naming the cluster that served the session's last request. +pub const HEADER_STARGATE_LAST_CLUSTER_ID: &str = "x-stargate-last-cluster-id"; pub const RETRY_REASON_QUEUE_ESTIMATE_MISMATCH: &str = "queue_estimate_mismatch"; pub const RETRY_REASON_UPSTREAM_ADMISSION_REJECTED: &str = "upstream_admission_rejected"; @@ -70,6 +72,13 @@ mod tests { assert_eq!(HEADER_STARGATE_RETRYABLE, "x-stargate-retryable"); assert_eq!(HEADER_STARGATE_RETRY_REASON, "x-stargate-retry-reason"); assert_eq!(HEADER_STARGATE_RETRY_AFTER_MS, "x-stargate-retry-after-ms"); + assert_eq!( + HEADER_STARGATE_LAST_CLUSTER_ID, + "x-stargate-last-cluster-id" + ); + assert!(is_internal_control_header( + &http::header::HeaderName::from_static(HEADER_STARGATE_LAST_CLUSTER_ID) + )); assert_eq!(HEADER_INFERENCE_SERVER_ID, "x-inference-server-id"); assert_eq!(HEADER_REVERSE_AUTH_TOKEN, "x-stargate-auth-token"); assert_eq!(WEBTRANSPORT_TUNNEL_PATH, "/_stargate/webtransport"); diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs index b14f5fff99..8ae7f54e28 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs @@ -24,7 +24,7 @@ use std::time::{Duration, Instant}; use anyhow::Context; use clap::ValueEnum; use stargate::load_balancer::{ - LoadBalancerAlgorithm, LoadBalancerAlgorithmConfig, LoadBalancerConfig, + LastClusterHint, LoadBalancerAlgorithm, LoadBalancerAlgorithmConfig, LoadBalancerConfig, LoadBalancerModelConfig, LoadBalancerRequest, LoadBalancerRouter, LoadBalancerTargetState, MAX_POWER_OF_N_SAMPLE_COUNT, }; @@ -36,6 +36,7 @@ enum WaitAndWidenTuning { IgnoreQueue, RttOnly, Affinity, + LastCluster, } struct LbMicrobenchScenarioMetadata { @@ -45,7 +46,7 @@ struct LbMicrobenchScenarioMetadata { excluded_clusters: usize, } -use WaitAndWidenTuning::{Affinity, IgnoreQueue, RttOnly}; +use WaitAndWidenTuning::{Affinity, IgnoreQueue, LastCluster, RttOnly}; macro_rules! scenarios { ($($scenario:ident, $model_id:literal, $algorithm:ident, $tuning:expr, $excluded:literal;)+) => { @@ -54,7 +55,7 @@ macro_rules! scenarios { #[value(rename_all = "kebab-case")] pub enum LbMicrobenchScenario { $($scenario,)+ } - const LB_MICROBENCH_SCENARIOS: [LbMicrobenchScenarioMetadata; 23] = [$( + const LB_MICROBENCH_SCENARIOS: [LbMicrobenchScenarioMetadata; 25] = [$( LbMicrobenchScenarioMetadata { model_id: $model_id, algorithm: LoadBalancerAlgorithm::$algorithm, @@ -87,6 +88,8 @@ scenarios! { WaitAndWidenAffinityMultiExcluded, "lb-bench-wait-and-widen-affinity-multi-excluded", WaitAndWiden, Some(Affinity), 2; Pulsar, "lb-bench-pulsar", Pulsar, None, 0; PulsarOneExcluded, "lb-bench-pulsar-one-excluded", Pulsar, None, 1; + PulsarWaitAndWiden, "lb-bench-pulsar-wait-and-widen", PulsarWaitAndWiden, None, 0; + PulsarWaitAndWidenLastCluster, "lb-bench-pulsar-wait-and-widen-last-cluster", PulsarWaitAndWiden, Some(LastCluster), 0; Random, "lb-bench-random", Random, None, 0; RandomOneExcluded, "lb-bench-random-one-excluded", Random, None, 1; RoundRobinOneExcluded, "lb-bench-round-robin-one-excluded", RoundRobin, None, 1; @@ -247,6 +250,17 @@ fn run_scenario( model_id: scenario.metadata().model_id.to_string(), }; let received_at = Instant::now(); + let run = LbMicrobenchRun { + target: &target, + target_state: &target_state, + candidates, + cache_keys, + config, + received_at, + excluded_cluster_ids: excluded_cluster_ids.as_ref(), + last_cluster_hints: scenario.metadata().tuning == Some(LastCluster), + }; + let warmup_hints = last_cluster_hints(&router, run, 0..config.warmup_iterations); for iteration in 0..config.warmup_iterations { let request = request_for_iteration( &target, @@ -254,23 +268,13 @@ fn run_scenario( iteration, received_at, excluded_cluster_ids.as_ref(), + warmup_hints.get(iteration), ); let _ = black_box(router.choose_candidate(&target_state, &request, candidates)); } let worker_count = config.concurrency.min(config.iterations); - let (stats, total_ns) = run_concurrent_measured_iterations( - &router, - LbMicrobenchRun { - target: &target, - target_state: &target_state, - candidates, - cache_keys, - config, - received_at, - excluded_cluster_ids: excluded_cluster_ids.as_ref(), - }, - ); + let (stats, total_ns) = run_concurrent_measured_iterations(&router, run); let choices = stats.candidate_counts.iter().sum::(); let avg_rank_depth = if choices == 0 { 0.0 @@ -341,6 +345,38 @@ struct LbMicrobenchRun<'a> { config: &'a LbMicrobenchConfig, received_at: Instant, excluded_cluster_ids: Option<&'a HashSet>, + /// Send a returning-session hint that names a cluster other than the + /// request's no-hint choice, so the decision promotes it. + last_cluster_hints: bool, +} + +/// One fresh hint per iteration, built outside the timed region. A hint pins +/// its classification on first use, so hints are never reused. +fn last_cluster_hints( + router: &LoadBalancerRouter, + run: LbMicrobenchRun<'_>, + iterations: std::ops::Range, +) -> Vec { + if !run.last_cluster_hints { + return Vec::new(); + } + iterations + .map(|iteration| { + let request = request_for_iteration( + run.target, + run.cache_keys, + iteration, + run.received_at, + run.excluded_cluster_ids, + None, + ); + let first_choice = router + .choose_candidate(run.target_state, &request, run.candidates) + .map_or(0, |choice| choice.candidate_index); + let last_index = (first_choice + 1) % run.candidates.len(); + LastClusterHint::new(Some(run.candidates[last_index].cluster_id.clone())) + }) + .collect() } fn run_concurrent_measured_iterations( @@ -360,16 +396,23 @@ fn run_concurrent_measured_iterations( handles.push(scope.spawn(move || { let measurements = LbMicrobenchMeasurements::new(run.candidates.len(), iterations.len()); + let first_iteration = run.config.warmup_iterations + iterations.start; + let hints = last_cluster_hints( + router, + run, + first_iteration..first_iteration + iterations.len(), + ); ready_barrier.wait(); release_barrier.wait(); let mut measurements = measurements; - for measured_iteration in iterations { + for (offset, measured_iteration) in iterations.enumerate() { let request = request_for_iteration( run.target, run.cache_keys, run.config.warmup_iterations + measured_iteration, run.received_at, run.excluded_cluster_ids, + hints.get(offset), ); if let Some(choice) = black_box(router.choose_candidate( run.target_state, @@ -449,6 +492,7 @@ fn request_for_iteration<'a>( iteration: usize, received_at: Instant, excluded_cluster_ids: Option<&'a HashSet>, + last_cluster: Option<&'a LastClusterHint>, ) -> LoadBalancerRequest<'a> { LoadBalancerRequest { routing_target: target, @@ -458,6 +502,7 @@ fn request_for_iteration<'a>( received_at, request_slo: Some(Duration::from_millis(250)), excluded_cluster_ids, + last_cluster, } } @@ -521,6 +566,9 @@ fn config_for_scenario( if metadata.tuning == Some(RttOnly) { wait_and_widen.ignore_input_processing_time = Some(true); } + if metadata.tuning == Some(LastCluster) { + wait_and_widen.last_cluster_affinity = Some(true); + } } config } @@ -608,7 +656,7 @@ mod tests { config.scenarios.clear(); let rows = run_lb_microbench(&config).expect("microbench should run"); - assert_eq!(rows.len(), 23); + assert_eq!(rows.len(), 25); assert_eq!(rows[0].scenario, LbMicrobenchScenario::PowerOfN); for row in rows { assert_eq!(row.choices, row.iterations); @@ -766,7 +814,7 @@ mod tests { for measured_iteration in 0..2 { let iteration = 2 + measured_iteration; assert_eq!( - request_for_iteration(&target, &cache_keys, iteration, received_at, None) + request_for_iteration(&target, &cache_keys, iteration, received_at, None, None) .cache_affinity_key, Some(cache_keys[2 + measured_iteration].as_str()) ); diff --git a/src/libraries/rust/stargate/crates/stargate-routing-sim/src/sim.rs b/src/libraries/rust/stargate/crates/stargate-routing-sim/src/sim.rs index 3ef45ec40a..65a2473e33 100644 --- a/src/libraries/rust/stargate/crates/stargate-routing-sim/src/sim.rs +++ b/src/libraries/rust/stargate/crates/stargate-routing-sim/src/sim.rs @@ -418,6 +418,7 @@ impl Simulation<'_> { received_at, request_slo: self.config.client.request_slo_ms.map(Duration::from_millis), excluded_cluster_ids: (!excluded.is_empty()).then_some(excluded), + last_cluster: None, }; let decision = self.stargates[stargate] .load_balancer diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs index 595d979656..0e7471f00a 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs @@ -558,6 +558,7 @@ mod tests { .resolve_algorithm_override(&target.model_id, None) .unwrap(), request_inputs: ProxyRequestInputs { + request_id: "req-1".to_string(), target, input_tokens: 1, priority: 0, @@ -566,6 +567,7 @@ mod tests { cache_affinity_key: None, routing_algorithm_override: None, routing_expression: None, + last_cluster: Default::default(), }, endpoint_name: "chat_completions", method: Method::POST, diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/request.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/request.rs index 77fd800a26..b404e6cab9 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/request.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/request.rs @@ -16,12 +16,15 @@ use axum::http::{HeaderMap, StatusCode}; use stargate_protocol::tunnel_contract::{ HEADER_INPUT_TOKENS, HEADER_MODEL, HEADER_PRIORITY, HEADER_REQUEST_ID, HEADER_ROUTING_KEY, + HEADER_STARGATE_LAST_CLUSTER_ID, }; use tracing::{Span, warn}; use super::ProxyRequestError; use crate::load_balancer::expression::RejectionError; -use crate::load_balancer::{LoadBalancerAlgorithmConfig, LoadBalancerAlgorithmOverride}; +use crate::load_balancer::{ + LastClusterHint, LoadBalancerAlgorithmConfig, LoadBalancerAlgorithmOverride, +}; use crate::routing_state::RoutingTargetKey; use super::{ @@ -30,6 +33,7 @@ use super::{ #[derive(Clone, Debug, PartialEq, Eq)] pub(super) struct ProxyRequestInputs { + pub(super) request_id: String, pub(super) target: RoutingTargetKey, pub(super) input_tokens: u64, pub(super) priority: u32, @@ -38,6 +42,8 @@ pub(super) struct ProxyRequestInputs { pub(super) cache_affinity_key: Option, pub(super) routing_algorithm_override: Option, pub(super) routing_expression: Option, + /// Consumed here; `x-stargate-*` headers are never forwarded upstream. + pub(super) last_cluster: LastClusterHint, } pub(super) fn parse_proxy_request_inputs( @@ -49,7 +55,8 @@ pub(super) fn parse_proxy_request_inputs( value.to_str().unwrap_or(""), ); } - get_optional_header(headers, HEADER_REQUEST_ID).ok_or(StatusCode::BAD_REQUEST)?; + let request_id = + get_optional_header(headers, HEADER_REQUEST_ID).ok_or(StatusCode::BAD_REQUEST)?; let input_tokens = parse_optional_numeric_header(headers, HEADER_INPUT_TOKENS)? .ok_or(StatusCode::BAD_REQUEST)?; let target = RoutingTargetKey::new( @@ -58,6 +65,7 @@ pub(super) fn parse_proxy_request_inputs( ); let routing_algorithm_override = parse_routing_algorithm_override(headers, &target)?; Ok(ProxyRequestInputs { + request_id, target, input_tokens, priority: parse_optional_numeric_header(headers, HEADER_PRIORITY)?.unwrap_or(0), @@ -70,6 +78,12 @@ pub(super) fn parse_proxy_request_inputs( .and_then(|value| value.to_str().ok()) .filter(|raw| raw.contains(';')) .map(ToOwned::to_owned), + // Blank or non-UTF-8 values read as absent, so a bad hint never fails + // the request. + last_cluster: LastClusterHint::new(get_optional_header( + headers, + HEADER_STARGATE_LAST_CLUSTER_ID, + )), }) } @@ -198,6 +212,44 @@ mod tests { ); } + #[test] + fn last_cluster_header_is_trimmed_and_request_id_is_kept() { + let mut headers = proxy_headers(); + set_header( + &mut headers, + HEADER_STARGATE_LAST_CLUSTER_ID, + " cluster-b ", + ); + + let inputs = parse_proxy_request_inputs(&headers).expect("headers should parse"); + + assert_eq!(inputs.request_id, "req-test"); + assert_eq!(inputs.last_cluster.last_cluster_id(), Some("cluster-b")); + assert_eq!(inputs.last_cluster.classification(), None); + } + + #[test] + fn unusable_last_cluster_header_reads_as_absent() { + for value in [ + HeaderValue::from_static(""), + HeaderValue::from_static(" "), + HeaderValue::from_bytes(b"\xff\xfe").unwrap(), + ] { + let mut headers = proxy_headers(); + headers.insert( + HeaderName::from_static(HEADER_STARGATE_LAST_CLUSTER_ID), + value, + ); + + let inputs = parse_proxy_request_inputs(&headers) + .expect("an unusable hint must not fail the request"); + assert_eq!(inputs.last_cluster.last_cluster_id(), None); + } + + let inputs = parse_proxy_request_inputs(&proxy_headers()).unwrap(); + assert_eq!(inputs.last_cluster.last_cluster_id(), None); + } + #[test] fn optional_u64_proxy_headers_reject_invalid_values() { for header in [ diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/routing.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/routing.rs index 82273fbad4..fc67f76415 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/routing.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/routing.rs @@ -339,6 +339,7 @@ mod tests { received_at: Instant::now(), request_slo: None, excluded_cluster_ids: None, + last_cluster: None, } } diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs index c1cb9332dc..177606679c 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs @@ -113,6 +113,7 @@ impl<'a> ProxyRequestRun<'a> { received_at: self.request.request_start, request_slo: request_inputs.request_slo_ms.map(Duration::from_millis), excluded_cluster_ids: self.excluded_cluster_ids(), + last_cluster: Some(&request_inputs.last_cluster), } } @@ -198,6 +199,16 @@ impl<'a> ProxyRequestRun<'a> { self.app.metrics.as_ref(), &self.request.request_inputs.target, ); + crate::load_balancer::record_session_selection( + self.app.metrics.as_ref(), + self.request.lb_resolution.config(), + &selected_cluster.routing_algorithm, + &self.request.request_inputs.request_id, + &self.request.request_inputs.target, + &self.request.request_inputs.last_cluster, + &selected_cluster.cluster.snapshot().cluster_id, + selected_cluster.selection.choice.rank_depth, + ); let comparator = self.request.lb_resolution.config().comparator(); loop { @@ -390,6 +401,7 @@ mod tests { fn request_inputs() -> ProxyRequestInputs { ProxyRequestInputs { + request_id: "req-1".to_string(), target: target(), input_tokens: 128, priority: 0, @@ -398,6 +410,7 @@ mod tests { cache_affinity_key: None, routing_algorithm_override: None, routing_expression: None, + last_cluster: Default::default(), } } diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs index 12fdc656f7..54e3cee216 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs @@ -64,6 +64,7 @@ pub(super) fn proxy_openai_request_span(headers: &HeaderMap) -> Span { routing.comparator = field::Empty, routing.rank_depth = field::Empty, routing.selected_after_kv_free_tokens_skip = field::Empty, + routing.session_state = field::Empty, routing.retry_attempts = field::Empty, routing.admission_rejection_reason = field::Empty, proxy.upstream_status = field::Empty, @@ -188,5 +189,89 @@ mod tests { assert!(fields.field("routing.sample_count_configured").is_some()); assert!(fields.field("routing.sample_count_effective").is_some()); assert!(fields.field("routing.comparator").is_some()); + assert!(fields.field("routing.session_state").is_some()); + } + + #[derive(Clone, Default)] + struct CapturedLogs(std::sync::Arc>>); + + impl std::io::Write for CapturedLogs { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + #[test] + fn session_selection_records_span_state_and_debug_event() { + use crate::load_balancer::{ + LastClusterHint, LoadBalancerAlgorithm, LoadBalancerAlgorithmConfig, + SessionClassification, SessionState, record_session_selection, + }; + use crate::metrics::StargateMetrics; + use crate::routing_state::RoutingTargetKey; + + let logs = CapturedLogs::default(); + let subscriber = tracing_subscriber::fmt() + .with_ansi(false) + .without_time() + .with_max_level(tracing::Level::DEBUG) + .with_writer({ + let logs = logs.clone(); + move || logs.clone() + }) + .finish(); + let metrics = StargateMetrics::new().unwrap(); + let config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PulsarWaitAndWiden); + let target = RoutingTargetKey::new(Some("tenant-a".to_string()), "model-a"); + let hint = LastClusterHint::classified( + Some("cluster-b"), + SessionClassification { + state: SessionState::Returning, + promoted: true, + }, + ); + + tracing::subscriber::with_default(subscriber, || { + let span = proxy_openai_request_span(&HeaderMap::new()); + let _entered = span.enter(); + record_session_selection( + &metrics, + &config, + "pulsar-wait-and-widen", + "req-42", + &target, + &hint, + "cluster-b", + 1, + ); + }); + + let output = String::from_utf8(logs.0.lock().unwrap().clone()).unwrap(); + let event = output + .lines() + .find(|line| line.contains("routing session classified")) + .unwrap_or_else(|| panic!("missing session debug event in {output}")); + assert!(event.contains("DEBUG"), "{event}"); + // The span prefix shows the recorded field value. + assert!( + event.contains("routing.session_state=\"returning\""), + "{event}" + ); + for field in [ + "request_id=req-42", + "model_id=model-a", + "routing_key=Some(\"tenant-a\")", + "session_state=\"returning\"", + "last_cluster_promoted=true", + "selected_cluster_id=cluster-b", + "selection=\"primary\"", + ] { + assert!(event.contains(field), "missing {field} in {event}"); + } } } diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/upstream.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/upstream.rs index 2adfe8a90b..b84a1a9c68 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/upstream.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/upstream.rs @@ -226,6 +226,7 @@ mod tests { for name in [ "X-Stargate-Upstream-Retryable", "X-Stargate-Auth-Token", + "X-Stargate-Last-Cluster-Id", "X-Stargate-Max-Wait-Ms", "X-Stargate-Additional-Control", ] { diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs index 2ce45992a4..999f70e30d 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs @@ -221,6 +221,11 @@ pub struct WaitAndWidenAlgorithmConfig { pub max_queued: Option, pub ignore_queue_time: Option, pub ignore_input_processing_time: Option, + /// Use `x-stargate-last-cluster-id` to tell new sessions from returning + /// ones. New sessions skip the affinity wait and prefill discount; + /// returning sessions put their last cluster first in the affinity order. + /// Defaults to `false`. + pub last_cluster_affinity: Option, } impl WaitAndWidenAlgorithmConfig { diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/expression.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/expression.rs index eaedb693eb..97db717948 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/expression.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/expression.rs @@ -628,6 +628,14 @@ mod tests { for (raw, class) in [ ("fastest;seed=x", "unknown_method"), ("pulsar;widen=2", "unknown_parameter"), + ( + "wait-and-widen;last_cluster_affinity=true", + "unknown_parameter", + ), + ( + "pulsar-wait-and-widen;last_cluster_affinity=true", + "unknown_parameter", + ), ("round-robin;seed=x", "not_applicable"), ("pulsar-wait-and-widen;comparator=ttft", "not_applicable"), ( diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs index 6433a4be78..627f42d65e 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs @@ -36,6 +36,7 @@ mod random; mod request; mod round_robin; mod router; +mod session; mod target_state; #[cfg(test)] mod tests; @@ -58,5 +59,7 @@ pub use request::{LoadBalancerCandidateChoice, LoadBalancerDecision, LoadBalance pub use router::{ LoadBalancerAlgorithmResolution, LoadBalancerCandidateSelection, LoadBalancerRouter, }; +pub(crate) use session::record_session_selection; +pub use session::{LastClusterHint, SessionClassification, SessionSelection, SessionState}; pub(crate) use target_state::LoadBalancerDefinition; pub use target_state::LoadBalancerTargetState; diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/power_of_n.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/power_of_n.rs index 4eaf1a3c17..2b90dcf6c3 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/power_of_n.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/power_of_n.rs @@ -240,6 +240,7 @@ mod tests { received_at: Instant::now(), request_slo: None, excluded_cluster_ids: Some(excluded_cluster_ids), + last_cluster: None, }; PowerOfNLoadBalancer { sample_count, @@ -264,6 +265,7 @@ mod tests { received_at: Instant::now(), request_slo: None, excluded_cluster_ids, + last_cluster: None, }; sample_candidates(&request, candidates, sample_count, rng) .as_slice() @@ -434,6 +436,7 @@ mod tests { received_at: Instant::now(), request_slo: None, excluded_cluster_ids: None, + last_cluster: None, }; let load_balancer = PowerOfNLoadBalancer { sample_count: 3, diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs index 7d4d6058b5..d66b996a77 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs @@ -13,9 +13,11 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::borrow::Cow; use std::time::Duration; use super::pulsar::PulsarLoadBalancer; +use super::session::{self, MissingLastCluster}; use super::wait_and_widen::{WaitAndWidenConfig, WaitAndWidenLoadBalancer}; use super::{ LoadBalancer, LoadBalancerAlgorithmConfig, LoadBalancerCandidateChoice, LoadBalancerDecision, @@ -39,15 +41,18 @@ pub(super) struct PulsarWaitAndWidenLoadBalancer { fallback_wait_and_widen: WaitAndWidenLoadBalancer, affinity_group_size: usize, affinity_wait: Duration, - band_widen_interval: Duration, + /// Explicit `band_widen_interval_ms`. Unset follows the request's + /// effective affinity wait. + band_widen_interval: Option, affinity_input_tokens_scale: f64, + last_cluster_affinity: bool, } /// Affinity selection discounts prefill and uses `max_queued`; fallback uses /// full prefill cost and `fallback_max_queued`. #[derive(Clone, Copy)] enum SelectionPhase { - Affinity, + Affinity { input_tokens_scale: f64 }, Fallback, } @@ -57,10 +62,7 @@ impl PulsarWaitAndWidenLoadBalancer { .wait_and_widen_settings() .expect("pulsar-wait-and-widen config has wait_and_widen settings"); let wait_and_widen_config = WaitAndWidenConfig::from_algorithm_config(&config)?; - let band_widen_interval = settings.band_widen_interval_ms.map_or( - wait_and_widen_config.cache_affinity_wait, - Duration::from_millis, - ); + let band_widen_interval = settings.band_widen_interval_ms.map(Duration::from_millis); let mut fallback_config = wait_and_widen_config.clone(); fallback_config.max_queued = settings.fallback_max_queued.unwrap_or(0); Ok(Self { @@ -71,6 +73,7 @@ impl PulsarWaitAndWidenLoadBalancer { affinity_wait: wait_and_widen_config.cache_affinity_wait, band_widen_interval, affinity_input_tokens_scale: wait_and_widen_config.cache_affinity_input_tokens_scale, + last_cluster_affinity: wait_and_widen_config.last_cluster_affinity, fallback_wait_and_widen: WaitAndWidenLoadBalancer::new(fallback_config), wait_and_widen: WaitAndWidenLoadBalancer::new(wait_and_widen_config), ranking: PulsarLoadBalancer::new(config), @@ -88,7 +91,9 @@ impl PulsarWaitAndWidenLoadBalancer { elapsed: Duration, ) -> LoadBalancerDecision { let (selector, input_tokens_scale) = match phase { - SelectionPhase::Affinity => (&self.wait_and_widen, self.affinity_input_tokens_scale), + SelectionPhase::Affinity { input_tokens_scale } => { + (&self.wait_and_widen, input_tokens_scale) + } SelectionPhase::Fallback => (&self.fallback_wait_and_widen, 1.0), }; let eligible = ranked_indices[..open_end] @@ -145,30 +150,35 @@ impl PulsarWaitAndWidenLoadBalancer { candidates: &[RoutedClusterSnapshot], elapsed: Duration, ) -> LoadBalancerDecision { - let ranked_indices = self.ranking.compute_ranking(request, candidates); - let Some(&primary_index) = ranked_indices.first() else { + let ranking = self.ranking.compute_ranking(request, candidates); + if ranking.is_empty() { return LoadBalancerDecision::Unavailable; - }; - if !self.wait_and_widen.has_queue_slo(request) - && self - .ranking - .feasibility(request, &candidates[primary_index]) - .is_eligible() - { - return LoadBalancerDecision::Selected(LoadBalancerCandidateChoice { - candidate_index: primary_index, - rank_depth: 1, - selected_after_kv_free_tokens_skip: false, - }); } + // Session rules reorder a per-request copy of the ranking. Rank 1 has + // no shortcut: it goes through queue admission, the affinity wait, and + // the affinity discount like every other affinity member. + let plan = session::plan_affinity( + self.last_cluster_affinity, + request, + candidates, + Cow::Owned(ranking), + MissingLastCluster::Keep, + self.affinity_wait, + self.affinity_input_tokens_scale, + ); + let ranked_indices = &*plan.order; + let affinity_wait = plan.affinity_wait; + let band_widen_interval = self.band_widen_interval.unwrap_or(affinity_wait); // Every attempt checks the affinity group first, even after the wait. let group_end = self.affinity_group_size.min(ranked_indices.len()); let affinity_bucket_wait = match self.decide_from_ranking_prefix( - SelectionPhase::Affinity, + SelectionPhase::Affinity { + input_tokens_scale: plan.input_tokens_scale, + }, request, candidates, - &ranked_indices, + ranked_indices, group_end, elapsed, ) { @@ -178,8 +188,8 @@ impl PulsarWaitAndWidenLoadBalancer { LoadBalancerDecision::Wait(delay) => Some(delay), LoadBalancerDecision::Unavailable => None, }; - if elapsed < self.affinity_wait { - let remaining = self.affinity_wait - elapsed; + if elapsed < affinity_wait { + let remaining = affinity_wait - elapsed; return LoadBalancerDecision::Wait( affinity_bucket_wait.map_or(remaining, |delay| delay.min(remaining)), ); @@ -190,11 +200,11 @@ impl PulsarWaitAndWidenLoadBalancer { // covers the whole ranking. Open candidates compete at full prefill // cost, as WaitAndWiden global buckets do, and bucket unlocks count // from the end of the affinity wait. - let widened_for = elapsed - self.affinity_wait; - let bands_open = if self.band_widen_interval.is_zero() { + let widened_for = elapsed - affinity_wait; + let bands_open = if band_widen_interval.is_zero() { usize::MAX } else { - usize::try_from(widened_for.as_nanos() / self.band_widen_interval.as_nanos() + 1) + usize::try_from(widened_for.as_nanos() / band_widen_interval.as_nanos() + 1) .unwrap_or(usize::MAX) }; let mut open_end = group_end; @@ -209,16 +219,15 @@ impl PulsarWaitAndWidenLoadBalancer { band_width = band_width.saturating_mul(2); } let next_band_wait = (open_end < ranked_indices.len()).then(|| { - let next_opening = self - .band_widen_interval - .saturating_mul(u32::try_from(bands_open).unwrap_or(u32::MAX)); + let next_opening = + band_widen_interval.saturating_mul(u32::try_from(bands_open).unwrap_or(u32::MAX)); next_opening.saturating_sub(widened_for) }); let open_decision = self.decide_from_ranking_prefix( SelectionPhase::Fallback, request, candidates, - &ranked_indices, + ranked_indices, open_end, widened_for, ); @@ -269,9 +278,12 @@ mod tests { use super::super::pulsar::PulsarLoadBalancer; use super::super::tests::LoadBalancerTestChoiceExt; + use std::borrow::Cow; + + use super::super::session::{self, MissingLastCluster}; use super::super::{ - LoadBalancerAlgorithm, LoadBalancerAlgorithmConfig, LoadBalancerRequest, - WaitAndWidenAlgorithmConfig, + LastClusterHint, LoadBalancerAlgorithm, LoadBalancerAlgorithmConfig, LoadBalancerRequest, + SessionClassification, SessionState, WaitAndWidenAlgorithmConfig, }; use super::*; use crate::routing_state::{RoutedClusterSnapshot, RoutingTargetKey}; @@ -293,6 +305,7 @@ mod tests { received_at: Instant::now(), request_slo: None, excluded_cluster_ids: None, + last_cluster: None, } } @@ -385,36 +398,87 @@ mod tests { ranked_indices } - #[test] - fn without_queue_slo_keeps_full_primary() { + /// No queue-SLO fields, rank 1 at its engine concurrency limit. + fn no_slo_full_primary( + seed: &str, + affinity_key: &str, + configure: impl FnOnce(&mut WaitAndWidenAlgorithmConfig), + ) -> ( + PulsarWaitAndWidenLoadBalancer, + Vec, + Vec, + ) { let target = target(); - let affinity_key = "hybrid-prefix-without-slo"; - let mut candidates = candidates(3, 0); - for candidate in &mut candidates { - candidate.stats.max_engine_concurrency = 1; - } - let hybrid_request = request(&target, Some(affinity_key), Some(0)); - let primary_index = - pulsar_ranked_indices("hybrid-seed", &target, affinity_key, 0, &candidates)[0]; - candidates[primary_index].stats.num_running_queries = 1; - + let mut candidates = single_slot_candidates(3); + let ranked = pulsar_ranked_indices(seed, &target, affinity_key, 0, &candidates); + candidates[ranked[0]].stats.num_running_queries = 1; + candidates[ranked[2]].stats.num_running_queries = 1; let hybrid = PulsarWaitAndWidenLoadBalancer::new(pulsar_wait_and_widen_algorithm_config( - "hybrid-seed", + seed, false, |settings| { settings.n = Some(2); + configure(settings); }, )) .expect("valid hybrid config"); + (hybrid, candidates, ranked) + } + + #[test] + fn without_queue_slo_overflows_full_primary() { + let target = target(); + let affinity_key = "hybrid-prefix-without-slo"; + let (hybrid, candidates, ranked) = no_slo_full_primary("hybrid-seed", affinity_key, |_| {}); + let request = request(&target, Some(affinity_key), Some(0)); + + let choice = hybrid + .decide_at(&request, &candidates, Duration::ZERO) + .selected() + .expect("a free rank-2 cluster should take the overflow"); + assert_eq!(choice.candidate_index, ranked[1]); + assert_eq!(choice.rank_depth, 2); + } + + #[test] + fn without_queue_slo_queues_on_primary_within_max_queued() { + let target = target(); + let affinity_key = "hybrid-prefix-without-slo"; + let (hybrid, mut candidates, ranked) = + no_slo_full_primary("hybrid-seed", affinity_key, |settings| { + settings.max_queued = Some(4); + }); + // One running query and one queued. + candidates[ranked[0]].stats.num_running_queries = 2; + let request = request(&target, Some(affinity_key), Some(0)); + + let choice = hybrid + .decide_at(&request, &candidates, Duration::ZERO) + .selected() + .expect("rank 1 should admit the request into its queue"); + assert_eq!(choice.candidate_index, ranked[0]); + assert_eq!(choice.rank_depth, 1); + } + + #[test] + fn without_queue_slo_waits_for_affinity_then_widens() { + let target = target(); + let affinity_key = "hybrid-prefix-without-slo"; + let (hybrid, candidates, ranked) = + no_slo_full_primary("hybrid-seed", affinity_key, |settings| { + settings.cache_affinity_wait_ms = Some(300); + }); + let request = request(&target, Some(affinity_key), Some(0)); - let chosen = hybrid - .choose_for_test(&hybrid_request, &candidates) - .expect("no-SLO primary should remain selectable despite live capacity"); assert_eq!( - chosen.candidate.cluster_id, - candidates[primary_index].cluster_id + hybrid.decide_at(&request, &candidates, Duration::ZERO), + LoadBalancerDecision::Wait(Duration::from_millis(300)) ); - assert_eq!(chosen.rank_depth, 1); + let choice = hybrid + .decide_at(&request, &candidates, Duration::from_millis(300)) + .selected() + .expect("the request should widen after the affinity wait"); + assert_eq!(choice.candidate_index, ranked[1]); } #[test] @@ -818,4 +882,538 @@ mod tests { "free-slot overflow rank {widened}" ); } + + // Session-aware pulsar-wait-and-widen (`last_cluster_affinity`). + + const SESSION_WAIT: Duration = Duration::from_millis(300); + const SESSION_KEY: &str = "hybrid-session"; + + fn session_hybrid( + seed: &str, + last_cluster_affinity: bool, + configure: impl FnOnce(&mut WaitAndWidenAlgorithmConfig), + ) -> PulsarWaitAndWidenLoadBalancer { + PulsarWaitAndWidenLoadBalancer::new(pulsar_wait_and_widen_algorithm_config( + seed, + false, + |settings| { + settings.n = Some(2); + settings.cache_affinity_wait_ms = Some(300); + settings.cache_affinity_input_tokens_scale = Some(0.1); + settings.last_cluster_affinity = Some(last_cluster_affinity); + configure(settings); + }, + )) + .expect("valid hybrid config") + } + + fn session_request<'a>( + target: &'a RoutingTargetKey, + input_tokens: u64, + hint: &'a LastClusterHint, + ) -> LoadBalancerRequest<'a> { + LoadBalancerRequest { + last_cluster: Some(hint), + ..request(target, Some(SESSION_KEY), Some(input_tokens)) + } + } + + fn last_cluster(candidate: &RoutedClusterSnapshot) -> LastClusterHint { + LastClusterHint::new(Some(candidate.cluster_id.clone())) + } + + fn selected_index(decision: LoadBalancerDecision) -> usize { + match decision { + LoadBalancerDecision::Selected(choice) => choice.candidate_index, + other => panic!("expected a selection, got {other:?}"), + } + } + + fn session_ranking(seed: &str, candidates: &[RoutedClusterSnapshot]) -> Vec { + pulsar_ranked_indices(seed, &target(), SESSION_KEY, 0, candidates) + } + + #[test] + fn new_session_overflows_without_waiting() { + let target = target(); + let mut candidates = single_slot_candidates(3); + let ranked = session_ranking("new-seed", &candidates); + candidates[ranked[0]].stats.num_running_queries = 1; + candidates[ranked[2]].stats.num_running_queries = 1; + + for hint in [ + LastClusterHint::default(), + LastClusterHint::new(Some("gone".to_string())), + ] { + let hybrid = session_hybrid("new-seed", true, |_| {}); + let choice = hybrid + .decide_at( + &session_request(&target, 0, &hint), + &candidates, + Duration::ZERO, + ) + .selected() + .expect("a new session should overflow at once"); + assert_eq!(choice.candidate_index, ranked[1]); + } + + let flag_off = session_hybrid("new-seed", false, |_| {}); + let hint = LastClusterHint::default(); + assert_eq!( + flag_off.decide_at( + &session_request(&target, 0, &hint), + &candidates, + Duration::ZERO + ), + LoadBalancerDecision::Wait(SESSION_WAIT) + ); + } + + #[test] + fn new_session_matches_flag_off_without_wait_or_discount() { + let target = target(); + let all_free = single_slot_candidates(5); + let ranked = session_ranking("match-seed", &all_free); + let mut all_full = all_free.clone(); + for candidate in &mut all_full { + candidate.stats.num_running_queries = 1; + } + let mut only_second_free = all_full.clone(); + only_second_free[ranked[1]].stats.num_running_queries = 0; + let mut only_last_free = all_full.clone(); + only_last_free[ranked[4]].stats.num_running_queries = 0; + let baseline = session_hybrid("match-seed", false, |settings| { + settings.cache_affinity_wait_ms = Some(0); + settings.cache_affinity_input_tokens_scale = Some(1.0); + }); + + for candidates in [&all_free, &all_full, &only_second_free, &only_last_free] { + for elapsed_ms in [0, 50, 300, 1_000] { + let elapsed = Duration::from_millis(elapsed_ms); + for hint in [ + LastClusterHint::default(), + LastClusterHint::new(Some("gone".to_string())), + ] { + let hybrid = session_hybrid("match-seed", true, |_| {}); + let request = session_request(&target, 0, &hint); + assert_eq!( + hybrid.decide_at(&request, candidates, elapsed), + baseline.decide_at(&request, candidates, elapsed), + "at {elapsed_ms} ms" + ); + } + } + } + } + + #[test] + fn new_session_uses_full_prefill_in_affinity_group() { + let target = target(); + let mut candidates = single_slot_candidates(4); + let ranked = session_ranking("prefill-seed", &candidates); + // Full prefill favors the faster processor (1500 vs 2001 ms). + // Discounted prefill favors the closer backend (600 vs 201 ms). + candidates[ranked[0]].rtt = Duration::from_millis(500); + candidates[ranked[0]].stats.last_mean_input_tps = 1_000.0; + candidates[ranked[1]].rtt = Duration::from_millis(1); + candidates[ranked[1]].stats.last_mean_input_tps = 500.0; + assert_eq!(session_ranking("prefill-seed", &candidates), ranked); + let group_of_two = |settings: &mut WaitAndWidenAlgorithmConfig| { + settings.cache_affinity_backend_selection_count = Some(2); + }; + + let hint = LastClusterHint::default(); + let request = session_request(&target, 1_000, &hint); + let flag_off = session_hybrid("prefill-seed", false, group_of_two); + assert_eq!( + selected_index(flag_off.decide_at(&request, &candidates, Duration::ZERO)), + ranked[1] + ); + for hint in [ + LastClusterHint::default(), + LastClusterHint::new(Some("gone".to_string())), + ] { + let hybrid = session_hybrid("prefill-seed", true, group_of_two); + assert_eq!( + selected_index(hybrid.decide_at( + &session_request(&target, 1_000, &hint), + &candidates, + Duration::ZERO + )), + ranked[0] + ); + } + } + + #[test] + fn returning_session_waits_for_last_cluster() { + let target = target(); + let mut candidates = single_slot_candidates(3); + let ranked = session_ranking("returning-seed", &candidates); + candidates[ranked[1]].stats.num_running_queries = 1; + let hint = last_cluster(&candidates[ranked[1]]); + let hybrid = session_hybrid("returning-seed", true, |_| {}); + let request = session_request(&target, 0, &hint); + + assert_eq!( + hybrid.decide_at(&request, &candidates, Duration::ZERO), + LoadBalancerDecision::Wait(SESSION_WAIT) + ); + for elapsed_ms in (0..300).step_by(10) { + assert!( + matches!( + hybrid.decide_at(&request, &candidates, Duration::from_millis(elapsed_ms)), + LoadBalancerDecision::Wait(_) + ), + "returning session must wait for its last cluster at {elapsed_ms} ms" + ); + } + assert_eq!( + hint.classification(), + Some(SessionClassification { + state: SessionState::Returning, + promoted: true, + }) + ); + + candidates[ranked[1]].stats.num_running_queries = 0; + let choice = hybrid + .decide_at(&request, &candidates, Duration::ZERO) + .selected() + .unwrap(); + assert_eq!(choice.candidate_index, ranked[1]); + assert_eq!(choice.rank_depth, 1); + } + + #[test] + fn returning_session_discounts_only_the_last_cluster() { + let target = target(); + let mut candidates = single_slot_candidates(3); + let ranked = session_ranking("discount-seed", &candidates); + let (a, b, c) = (ranked[0], ranked[1], ranked[2]); + candidates[b].stats.num_running_queries = 1; + // A wins only when discounted (201 ms vs C at 1500 ms). + candidates[a].rtt = Duration::from_millis(1); + candidates[a].stats.last_mean_input_tps = 500.0; + candidates[c].rtt = Duration::from_millis(500); + candidates[c].stats.last_mean_input_tps = 1_000.0; + assert_eq!(session_ranking("discount-seed", &candidates), ranked); + let hint = last_cluster(&candidates[b]); + let request = session_request(&target, 1_000, &hint); + + let hybrid = session_hybrid("discount-seed", true, |_| {}); + assert_eq!( + selected_index(hybrid.decide_at(&request, &candidates, SESSION_WAIT)), + c + ); + let flag_off = session_hybrid("discount-seed", false, |_| {}); + assert_eq!( + selected_index(flag_off.decide_at(&request, &candidates, Duration::ZERO)), + a + ); + } + + #[test] + fn returning_session_outside_group_replaces_last_member() { + let target = target(); + let mut candidates = single_slot_candidates(4); + let ranked = session_ranking("group-seed", &candidates); + let (a, b, c, d) = (ranked[0], ranked[1], ranked[2], ranked[3]); + let hint = last_cluster(&candidates[c]); + let hybrid = session_hybrid("group-seed", true, |settings| { + settings.cache_affinity_backend_selection_count = Some(2); + }); + let request = session_request(&target, 0, &hint); + + // The group is (C, A): B is not selectable during the wait. + for _ in 0..32 { + let choice = selected_index(hybrid.decide_at(&request, &candidates, Duration::ZERO)); + assert!([a, c].contains(&choice), "chose {choice}"); + } + + for index in [a, c, d] { + candidates[index].stats.num_running_queries = 1; + } + assert_eq!( + hybrid.decide_at(&request, &candidates, Duration::from_millis(100)), + LoadBalancerDecision::Wait(Duration::from_millis(200)) + ); + let choice = hybrid + .decide_at(&request, &candidates, SESSION_WAIT) + .selected() + .unwrap(); + assert_eq!(choice.candidate_index, b); + assert_eq!(choice.rank_depth, 3); + } + + #[test] + fn returning_session_inside_group_keeps_members() { + let target = target(); + let mut candidates = single_slot_candidates(4); + let ranked = session_ranking("members-seed", &candidates); + let hint = last_cluster(&candidates[ranked[1]]); + let group_of_two = |settings: &mut WaitAndWidenAlgorithmConfig| { + settings.cache_affinity_backend_selection_count = Some(2); + }; + let hybrid = session_hybrid("members-seed", true, group_of_two); + let flag_off = session_hybrid("members-seed", false, group_of_two); + let request = session_request(&target, 0, &hint); + + for _ in 0..32 { + let choice = selected_index(hybrid.decide_at(&request, &candidates, Duration::ZERO)); + assert!(ranked[..2].contains(&choice), "chose {choice}"); + } + candidates[ranked[0]].stats.num_running_queries = 1; + candidates[ranked[1]].stats.num_running_queries = 1; + for elapsed_ms in [0, 100] { + let elapsed = Duration::from_millis(elapsed_ms); + assert_eq!( + hybrid.decide_at(&request, &candidates, elapsed), + flag_off.decide_at(&request, &candidates, elapsed) + ); + } + } + + #[test] + fn returning_session_widens_bands_from_reordered_ranking() { + let target = target(); + let base = single_slot_candidates(5); + let ranked = session_ranking("band-seed", &base); + let hint = last_cluster(&base[ranked[3]]); + let hybrid = session_hybrid("band-seed", true, |settings| { + settings.band_widen_interval_ms = Some(100); + }); + let request = session_request(&target, 0, &hint); + let only_free = |rank: usize| { + let mut candidates = base.clone(); + for (position, index) in ranked.iter().enumerate() { + if position != rank { + candidates[*index].stats.num_running_queries = 1; + } + } + candidates + }; + + // After the wait the open set is D, A, B. C opens with the next band. + let c_free = only_free(2); + assert_eq!( + hybrid.decide_at(&request, &c_free, SESSION_WAIT), + LoadBalancerDecision::Wait(Duration::from_millis(100)) + ); + let choice = hybrid + .decide_at(&request, &c_free, SESSION_WAIT + Duration::from_millis(100)) + .selected() + .unwrap(); + assert_eq!(choice.candidate_index, ranked[2]); + assert_eq!(choice.rank_depth, 4); + + let b_free = only_free(1); + let choice = hybrid + .decide_at(&request, &b_free, SESSION_WAIT) + .selected() + .unwrap(); + assert_eq!(choice.candidate_index, ranked[1]); + assert_eq!(choice.rank_depth, 3); + } + + #[test] + fn retry_excluded_last_cluster_acts_as_ineligible_primary() { + let target = target(); + let candidates = single_slot_candidates(3); + let ranked = session_ranking("excluded-seed", &candidates); + let b_id = candidates[ranked[1]].cluster_id.clone(); + let hint = last_cluster(&candidates[ranked[1]]); + let excluded = HashSet::from([b_id.clone()]); + let request = LoadBalancerRequest { + excluded_cluster_ids: Some(&excluded), + ..session_request(&target, 0, &hint) + }; + let hybrid = session_hybrid("excluded-seed", true, |_| {}); + + assert_eq!( + hybrid.decide_at(&request, &candidates, Duration::ZERO), + LoadBalancerDecision::Wait(SESSION_WAIT) + ); + let choice = selected_index(hybrid.decide_at(&request, &candidates, SESSION_WAIT)); + assert_ne!(candidates[choice].cluster_id, b_id); + } + + #[test] + fn kv_infeasible_last_cluster_acts_as_ineligible_primary() { + let target = target(); + let mut candidates = candidates(3, 1024); + let ranked = session_ranking("kv-seed", &candidates); + candidates[ranked[1]].stats.kv_cache_free_tokens = 50; + candidates[ranked[1]].stats.kv_cache_used_tokens = 974; + let hint = last_cluster(&candidates[ranked[1]]); + let request = session_request(&target, 100, &hint); + let hybrid = PulsarWaitAndWidenLoadBalancer::new(pulsar_wait_and_widen_algorithm_config( + "kv-seed", + true, + |settings| { + settings.n = Some(2); + settings.cache_affinity_wait_ms = Some(300); + settings.last_cluster_affinity = Some(true); + }, + )) + .expect("valid hybrid config"); + + assert_eq!( + hybrid.decide_at(&request, &candidates, Duration::ZERO), + LoadBalancerDecision::Wait(SESSION_WAIT) + ); + let choice = hybrid + .decide_at(&request, &candidates, SESSION_WAIT) + .selected() + .unwrap(); + assert_ne!(choice.candidate_index, ranked[1]); + assert!(choice.selected_after_kv_free_tokens_skip); + } + + #[test] + fn promotion_is_not_cached() { + let target = target(); + let candidates = single_slot_candidates(3); + let ranked = session_ranking("cache-seed", &candidates); + let hybrid = session_hybrid("cache-seed", true, |_| {}); + let returning = last_cluster(&candidates[ranked[1]]); + assert_eq!( + selected_index(hybrid.decide_at( + &session_request(&target, 0, &returning), + &candidates, + Duration::ZERO + )), + ranked[1] + ); + + let new_session = LastClusterHint::default(); + let choice = hybrid + .decide_at( + &session_request(&target, 0, &new_session), + &candidates, + Duration::ZERO, + ) + .selected() + .unwrap(); + assert_eq!(choice.candidate_index, ranked[0]); + assert_eq!(choice.rank_depth, 1); + } + + #[test] + fn session_decisions_match_across_instances() { + let target = target(); + let mut candidates = single_slot_candidates(5); + let ranked = session_ranking("instances-seed", &candidates); + let (a, b, c, d, e) = (ranked[0], ranked[1], ranked[2], ranked[3], ranked[4]); + // RTTs more than one TTFT bucket apart give one winner per open + // set; ties within a bucket break randomly. With D + // promoted the order is D, A, B, C, E: E is fastest but opens last. + for (rtt_ms, index) in [(150, a), (100, b), (50, c), (250, d), (0, e)] { + candidates[index].rtt = Duration::from_millis(rtt_ms); + } + let mut group_full = candidates.clone(); + group_full[a].stats.num_running_queries = 1; + group_full[d].stats.num_running_queries = 1; + // A 1 ms band interval opens E one millisecond after C's band, + // before any later TTFT bucket unlocks. + let configure = |settings: &mut WaitAndWidenAlgorithmConfig| { + settings.n = Some(5); + settings.cache_affinity_backend_selection_count = Some(2); + settings.band_widen_interval_ms = Some(1); + }; + let instances: Vec<_> = (0..2) + .map(|_| session_hybrid("instances-seed", true, configure)) + .collect(); + + let plans: Vec<_> = instances + .iter() + .map(|hybrid| { + let hint = last_cluster(&candidates[d]); + let request = session_request(&target, 0, &hint); + let plan = session::plan_affinity( + hybrid.last_cluster_affinity, + &request, + &candidates, + Cow::Owned(hybrid.ranking.compute_ranking(&request, &candidates)), + MissingLastCluster::Keep, + hybrid.affinity_wait, + hybrid.affinity_input_tokens_scale, + ); + ( + plan.order.into_owned(), + plan.affinity_wait, + plan.input_tokens_scale, + hint.classification(), + ) + }) + .collect(); + assert_eq!(plans[0], plans[1]); + assert_eq!(plans[0].0, vec![d, a, b, c, e]); + assert_eq!((plans[0].1, plans[0].2), (SESSION_WAIT, 0.1)); + + for (snapshot, elapsed_ms, expected) in [ + (&candidates, 0, Some(a)), + (&group_full, 150, None), + (&group_full, 300, Some(c)), + (&group_full, 301, Some(e)), + ] { + let elapsed = Duration::from_millis(elapsed_ms); + let decisions: Vec<_> = instances + .iter() + .map(|hybrid| { + let hint = last_cluster(&candidates[d]); + hybrid.decide_at(&session_request(&target, 0, &hint), snapshot, elapsed) + }) + .collect(); + assert_eq!(decisions[0], decisions[1], "at {elapsed_ms} ms"); + assert_eq!( + decisions[0].selected().map(|choice| choice.candidate_index), + expected, + "at {elapsed_ms} ms" + ); + } + } + + #[test] + fn last_cluster_is_ignored_when_flag_is_off() { + let target = target(); + let mut candidates = single_slot_candidates(3); + let ranked = session_ranking("off-seed", &candidates); + candidates[ranked[1]].stats.num_running_queries = 1; + let hybrid = session_hybrid("off-seed", false, |_| {}); + + for elapsed_ms in [0, 300] { + let elapsed = Duration::from_millis(elapsed_ms); + let hint = last_cluster(&candidates[ranked[1]]); + let with_header = + hybrid.decide_at(&session_request(&target, 0, &hint), &candidates, elapsed); + let without = hybrid.decide_at( + &request(&target, Some(SESSION_KEY), Some(0)), + &candidates, + elapsed, + ); + assert_eq!(with_header, without); + assert_eq!(hint.classification(), None); + } + } + + #[test] + fn last_cluster_needs_an_affinity_key() { + let target = target(); + let candidates = single_slot_candidates(3); + let hint = last_cluster(&candidates[0]); + let hybrid = session_hybrid("no-key-seed", true, |_| {}); + let request = LoadBalancerRequest { + last_cluster: Some(&hint), + ..request(&target, None, Some(0)) + }; + + assert!( + hybrid + .decide_at(&request, &candidates, Duration::ZERO) + .selected() + .is_some() + ); + assert_eq!(hint.classification(), None); + } } diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/random.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/random.rs index 9203803e75..047bf88ecd 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/random.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/random.rs @@ -137,6 +137,7 @@ mod tests { received_at: Instant::now(), request_slo: None, excluded_cluster_ids, + last_cluster: None, } } diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/request.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/request.rs index 5061859e96..040bd203b0 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/request.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/request.rs @@ -16,6 +16,7 @@ use std::collections::HashSet; use std::time::{Duration, Instant}; +use super::LastClusterHint; use crate::routing_state::RoutingTargetKey; #[derive(Clone, Debug)] @@ -27,6 +28,8 @@ pub struct LoadBalancerRequest<'a> { pub received_at: Instant, pub request_slo: Option, pub excluded_cluster_ids: Option<&'a HashSet>, + /// Last-cluster hint for session-aware affinity. `None` outside the proxy. + pub last_cluster: Option<&'a LastClusterHint>, } impl LoadBalancerRequest<'_> { diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/round_robin.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/round_robin.rs index a037e61dd6..ad6a8d0fcc 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/round_robin.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/round_robin.rs @@ -167,6 +167,7 @@ mod tests { received_at: Instant::now(), request_slo: None, excluded_cluster_ids, + last_cluster: None, } } diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/session.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/session.rs new file mode 100644 index 0000000000..7274586e6b --- /dev/null +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/session.rs @@ -0,0 +1,759 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Session-aware affinity for `last_cluster_affinity`. +//! +//! The gateway sends the cluster that served a session's last request in +//! `x-stargate-last-cluster-id`. A request without that hint is a new session: +//! it skips the affinity wait and the affinity prefill discount. A request +//! whose hint names a current candidate is a returning session: its last +//! cluster moves to the front of the affinity order, and the configured wait, +//! discount, and widening apply unchanged. + +use std::borrow::Cow; +use std::sync::OnceLock; +use std::time::Duration; + +use tracing::Span; + +use super::{LoadBalancerAlgorithmConfig, LoadBalancerRequest}; +use crate::metrics::StargateMetrics; +use crate::routing_state::{RoutedClusterSnapshot, RoutingTargetKey}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SessionState { + /// No usable last-cluster hint. + New, + /// The hint names a current candidate. + Returning, + /// The hint names a cluster outside the candidate set. Routed as `New`. + Stale, +} + +impl SessionState { + pub const ALL: [Self; 3] = [Self::New, Self::Returning, Self::Stale]; + + pub fn as_str(self) -> &'static str { + match self { + Self::New => "new", + Self::Returning => "returning", + Self::Stale => "stale", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SessionSelection { + Primary, + Fallback, +} + +impl SessionSelection { + pub const ALL: [Self; 2] = [Self::Primary, Self::Fallback]; + + pub fn as_str(self) -> &'static str { + match self { + Self::Primary => "primary", + Self::Fallback => "fallback", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct SessionClassification { + pub state: SessionState, + /// The last cluster moved ahead of the original first affinity position. + pub promoted: bool, +} + +/// Per-request last-cluster hint. The first load-balancer decision classifies +/// the request; retries and timed waits reuse that classification. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct LastClusterHint { + last_cluster_id: Option, + classification: OnceLock, +} + +impl LastClusterHint { + pub fn new(last_cluster_id: Option) -> Self { + Self { + last_cluster_id, + classification: OnceLock::new(), + } + } + + pub fn last_cluster_id(&self) -> Option<&str> { + self.last_cluster_id.as_deref() + } + + /// `None` until a session-aware load balancer classifies the request. + pub fn classification(&self) -> Option { + self.classification.get().copied() + } + + #[cfg(test)] + pub(crate) fn classified( + last_cluster_id: Option<&str>, + classification: SessionClassification, + ) -> Self { + let hint = Self::new(last_cluster_id.map(ToOwned::to_owned)); + hint.classification + .set(classification) + .expect("a new hint is unclassified"); + hint + } +} + +/// Affinity inputs for one decision after the session rules apply. +#[derive(Debug)] +pub(super) struct AffinityPlan<'a> { + pub(super) order: Cow<'a, [usize]>, + pub(super) affinity_wait: Duration, + pub(super) input_tokens_scale: f64, +} + +/// What to do when a returning session's last cluster is a candidate but not +/// in the affinity order. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum MissingLastCluster { + /// The order is a fixed-size affinity group: put the last cluster first + /// and drop the group's last member. + InsertAndTruncate, + /// The order is a full ranking that omits the cluster: leave it unchanged. + Keep, +} + +/// Applies the session rules to one decision. The order passed in is the +/// original affinity order; per-key caches are never modified. +pub(super) fn plan_affinity<'a>( + enabled: bool, + request: &LoadBalancerRequest<'_>, + candidates: &[RoutedClusterSnapshot], + order: Cow<'a, [usize]>, + missing: MissingLastCluster, + configured_wait: Duration, + configured_scale: f64, +) -> AffinityPlan<'a> { + if !enabled || request.cache_affinity_key.is_none() { + return AffinityPlan { + order, + affinity_wait: configured_wait, + input_tokens_scale: configured_scale, + }; + } + let hint_id = request + .last_cluster + .and_then(LastClusterHint::last_cluster_id); + let last_index = hint_id.and_then(|id| { + candidates + .iter() + .position(|candidate| candidate.cluster_id == id) + }); + let classify = || match (hint_id, last_index) { + (None, _) => SessionClassification { + state: SessionState::New, + promoted: false, + }, + (Some(_), None) => SessionClassification { + state: SessionState::Stale, + promoted: false, + }, + (Some(_), Some(last_index)) => SessionClassification { + state: SessionState::Returning, + promoted: order.first() != Some(&last_index) + && (order.contains(&last_index) + || missing == MissingLastCluster::InsertAndTruncate), + }, + }; + let classification = request + .last_cluster + .map_or_else(classify, |hint| *hint.classification.get_or_init(classify)); + + match classification.state { + SessionState::New | SessionState::Stale => AffinityPlan { + order, + affinity_wait: Duration::ZERO, + input_tokens_scale: 1.0, + }, + SessionState::Returning => AffinityPlan { + // A pinned returning session whose last cluster left the candidate + // set on a later attempt keeps the original order. + order: match last_index { + Some(last_index) => promote(order, last_index, missing), + None => order, + }, + affinity_wait: configured_wait, + input_tokens_scale: configured_scale, + }, + } +} + +fn promote( + mut order: Cow<'_, [usize]>, + last_index: usize, + missing: MissingLastCluster, +) -> Cow<'_, [usize]> { + match order.iter().position(|index| *index == last_index) { + Some(0) => {} + // Rotating the prefix keeps every other member's relative order. + Some(position) => order.to_mut()[..=position].rotate_right(1), + None if missing == MissingLastCluster::InsertAndTruncate && !order.is_empty() => { + let order = order.to_mut(); + order.pop(); + order.insert(0, last_index); + } + None => {} + } + order +} + +/// Records the session metric, span field, and debug log for a selection. +/// +/// For a returning session, `primary` means the request went to its last +/// cluster. For new and stale sessions, `primary` means the cluster came from +/// the affinity group (`rank_depth` within the configured group size). +#[allow(clippy::too_many_arguments)] +pub(crate) fn record_session_selection( + metrics: &StargateMetrics, + config: &LoadBalancerAlgorithmConfig, + algorithm: &str, + request_id: &str, + target: &RoutingTargetKey, + hint: &LastClusterHint, + selected_cluster_id: &str, + rank_depth: usize, +) { + let Some(classification) = hint.classification() else { + return; + }; + let selection = match classification.state { + SessionState::Returning => { + if hint.last_cluster_id() == Some(selected_cluster_id) { + SessionSelection::Primary + } else { + SessionSelection::Fallback + } + } + SessionState::New | SessionState::Stale => { + let group_size = config + .wait_and_widen_settings() + .and_then(|settings| settings.cache_affinity_backend_selection_count) + .filter(|count| *count > 0) + .unwrap_or(1); + if rank_depth <= group_size { + SessionSelection::Primary + } else { + SessionSelection::Fallback + } + } + }; + + let routing_key = target.routing_key.as_deref(); + // Routing key and model are dynamic, so every state and selection pair + // starts at zero the first time a target reports a session selection. + for state in SessionState::ALL { + for pair_selection in SessionSelection::ALL { + metrics + .routing_session_selections_total( + routing_key, + &target.model_id, + algorithm, + state.as_str(), + pair_selection.as_str(), + ) + .inc_by(0); + } + } + metrics + .routing_session_selections_total( + routing_key, + &target.model_id, + algorithm, + classification.state.as_str(), + selection.as_str(), + ) + .inc(); + + Span::current().record("routing.session_state", classification.state.as_str()); + tracing::debug!( + request_id = %request_id, + model_id = %target.model_id, + routing_key = ?target.routing_key, + algorithm = %algorithm, + session_state = classification.state.as_str(), + last_cluster_promoted = classification.promoted, + selected_cluster_id = %selected_cluster_id, + selection = selection.as_str(), + "routing session classified" + ); +} + +#[cfg(test)] +mod tests { + use std::time::Instant; + + use super::*; + use crate::load_balancer::tests::candidates; + use crate::load_balancer::{ + LoadBalancerAlgorithm, LoadBalancerAlgorithmSettings, WaitAndWidenAlgorithmConfig, + }; + + const WAIT: Duration = Duration::from_millis(300); + const SCALE: f64 = 0.1; + + fn target() -> RoutingTargetKey { + RoutingTargetKey::new(Some("rk-1".to_string()), "model-a") + } + + fn request<'a>( + target: &'a RoutingTargetKey, + cache_affinity_key: Option<&'a str>, + hint: Option<&'a LastClusterHint>, + ) -> LoadBalancerRequest<'a> { + LoadBalancerRequest { + routing_target: target, + cache_affinity_key, + input_tokens: Some(100), + priority: 0, + received_at: Instant::now(), + request_slo: None, + excluded_cluster_ids: None, + last_cluster: hint, + } + } + + fn plan<'a>( + enabled: bool, + request: &LoadBalancerRequest<'_>, + candidates: &[RoutedClusterSnapshot], + order: &'a [usize], + missing: MissingLastCluster, + ) -> AffinityPlan<'a> { + plan_affinity( + enabled, + request, + candidates, + Cow::Borrowed(order), + missing, + WAIT, + SCALE, + ) + } + + fn five_candidates() -> Vec { + candidates(&["A", "B", "C", "D", "E"]) + } + + fn hint(id: &str) -> LastClusterHint { + LastClusterHint::new(Some(id.to_string())) + } + + #[test] + fn missing_hint_classifies_new_and_drops_wait_and_discount() { + let target = target(); + let candidates = five_candidates(); + let hint = LastClusterHint::new(None); + let request = request(&target, Some("key"), Some(&hint)); + let plan = plan( + true, + &request, + &candidates, + &[0], + MissingLastCluster::InsertAndTruncate, + ); + + assert_eq!(&*plan.order, &[0]); + assert_eq!(plan.affinity_wait, Duration::ZERO); + assert_eq!(plan.input_tokens_scale, 1.0); + assert_eq!( + hint.classification(), + Some(SessionClassification { + state: SessionState::New, + promoted: false + }) + ); + } + + #[test] + fn unknown_cluster_classifies_stale_and_routes_as_new() { + let target = target(); + let candidates = five_candidates(); + let hint = hint("gone"); + let request = request(&target, Some("key"), Some(&hint)); + let plan = plan( + true, + &request, + &candidates, + &[0, 1], + MissingLastCluster::InsertAndTruncate, + ); + + assert_eq!(&*plan.order, &[0, 1]); + assert_eq!(plan.affinity_wait, Duration::ZERO); + assert_eq!(plan.input_tokens_scale, 1.0); + assert_eq!( + hint.classification().map(|c| c.state), + Some(SessionState::Stale) + ); + } + + #[test] + fn returning_session_promotes_into_single_member_group() { + let target = target(); + let candidates = five_candidates(); + let hint = hint("B"); + let request = request(&target, Some("key"), Some(&hint)); + let plan = plan( + true, + &request, + &candidates, + &[0], + MissingLastCluster::InsertAndTruncate, + ); + + assert_eq!(&*plan.order, &[1]); + assert_eq!(plan.affinity_wait, WAIT); + assert_eq!(plan.input_tokens_scale, SCALE); + assert_eq!( + hint.classification(), + Some(SessionClassification { + state: SessionState::Returning, + promoted: true + }) + ); + } + + #[test] + fn returning_session_outside_group_replaces_last_member() { + let target = target(); + let candidates = five_candidates(); + let hint = hint("C"); + let request = request(&target, Some("key"), Some(&hint)); + let plan = plan( + true, + &request, + &candidates, + &[0, 1], + MissingLastCluster::InsertAndTruncate, + ); + + assert_eq!(&*plan.order, &[2, 0]); + } + + #[test] + fn returning_session_inside_group_keeps_members() { + let target = target(); + let candidates = five_candidates(); + let hint = hint("B"); + let request = request(&target, Some("key"), Some(&hint)); + let plan = plan( + true, + &request, + &candidates, + &[0, 1], + MissingLastCluster::InsertAndTruncate, + ); + + assert_eq!(&*plan.order, &[1, 0]); + } + + #[test] + fn returning_session_reorders_full_ranking_and_keeps_relative_order() { + let target = target(); + let candidates = five_candidates(); + let hint = hint("D"); + let request = request(&target, Some("key"), Some(&hint)); + let plan = plan( + true, + &request, + &candidates, + &[0, 1, 2, 3, 4], + MissingLastCluster::Keep, + ); + + assert_eq!(&*plan.order, &[3, 0, 1, 2, 4]); + } + + #[test] + fn returning_session_already_first_is_not_promoted() { + let target = target(); + let candidates = five_candidates(); + let hint = hint("A"); + let request = request(&target, Some("key"), Some(&hint)); + let plan = plan( + true, + &request, + &candidates, + &[0, 1], + MissingLastCluster::Keep, + ); + + assert!(matches!(plan.order, Cow::Borrowed(_))); + assert_eq!(&*plan.order, &[0, 1]); + assert_eq!( + hint.classification(), + Some(SessionClassification { + state: SessionState::Returning, + promoted: false + }) + ); + } + + #[test] + fn ranking_without_last_cluster_is_kept() { + let target = target(); + let candidates = five_candidates(); + let hint = hint("E"); + let request = request(&target, Some("key"), Some(&hint)); + let plan = plan( + true, + &request, + &candidates, + &[0, 1, 2, 3], + MissingLastCluster::Keep, + ); + + assert_eq!(&*plan.order, &[0, 1, 2, 3]); + assert_eq!(plan.affinity_wait, WAIT); + assert_eq!( + hint.classification(), + Some(SessionClassification { + state: SessionState::Returning, + promoted: false + }) + ); + } + + #[test] + fn classification_is_pinned_across_attempts() { + let target = target(); + let candidates = five_candidates(); + let hint = hint("B"); + let request = request(&target, Some("key"), Some(&hint)); + let _ = plan( + true, + &request, + &candidates, + &[0], + MissingLastCluster::InsertAndTruncate, + ); + + // B left the candidate set before a retry: still returning, no promotion. + let later = crate::load_balancer::tests::candidates(&["A", "C"]); + let plan = plan( + true, + &request, + &later, + &[0], + MissingLastCluster::InsertAndTruncate, + ); + assert_eq!(&*plan.order, &[0]); + assert_eq!(plan.affinity_wait, WAIT); + assert_eq!(plan.input_tokens_scale, SCALE); + assert_eq!( + hint.classification().map(|c| c.state), + Some(SessionState::Returning) + ); + } + + #[test] + fn flag_off_or_missing_affinity_key_is_not_classified() { + let target = target(); + let candidates = five_candidates(); + for (enabled, affinity_key) in [(false, Some("key")), (true, None)] { + let hint = hint("B"); + let request = request(&target, affinity_key, Some(&hint)); + let plan = plan( + enabled, + &request, + &candidates, + &[0], + MissingLastCluster::InsertAndTruncate, + ); + + assert_eq!(&*plan.order, &[0]); + assert_eq!(plan.affinity_wait, WAIT); + assert_eq!(plan.input_tokens_scale, SCALE); + assert_eq!(hint.classification(), None); + } + } + + #[test] + fn request_without_hint_is_new_when_enabled() { + let target = target(); + let candidates = five_candidates(); + let request = request(&target, Some("key"), None); + let plan = plan( + true, + &request, + &candidates, + &[0], + MissingLastCluster::InsertAndTruncate, + ); + + assert_eq!(plan.affinity_wait, Duration::ZERO); + assert_eq!(plan.input_tokens_scale, 1.0); + } + + fn classified_hint(id: Option<&str>, state: SessionState) -> LastClusterHint { + LastClusterHint::classified( + id, + SessionClassification { + state, + promoted: false, + }, + ) + } + + fn session_series(metrics: &StargateMetrics) -> Vec<(String, String, u64)> { + metrics + .registry() + .gather() + .into_iter() + .filter(|family| family.name() == "stargate_routing_session_selections_total") + .flat_map(|family| family.get_metric().to_vec()) + .map(|metric| { + let label = |name: &str| { + metric + .get_label() + .iter() + .find(|label| label.name() == name) + .map(|label| label.value().to_string()) + .unwrap_or_default() + }; + ( + label("session_state"), + label("selection"), + metric.get_counter().value() as u64, + ) + }) + .collect() + } + + fn pulsar_wait_and_widen_config(group_size: Option) -> LoadBalancerAlgorithmConfig { + LoadBalancerAlgorithmConfig { + settings: LoadBalancerAlgorithmSettings::PulsarWaitAndWiden( + WaitAndWidenAlgorithmConfig { + cache_affinity_backend_selection_count: group_size, + last_cluster_affinity: Some(true), + ..WaitAndWidenAlgorithmConfig::default() + }, + ), + ..LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PulsarWaitAndWiden) + } + } + + #[test] + fn session_metric_starts_every_pair_at_zero_and_counts_selection() { + let metrics = StargateMetrics::new().unwrap(); + let config = pulsar_wait_and_widen_config(None); + let hint = classified_hint(Some("B"), SessionState::Returning); + record_session_selection( + &metrics, + &config, + "pulsar-wait-and-widen", + "req-1", + &target(), + &hint, + "B", + 2, + ); + + let mut series = session_series(&metrics); + series.sort(); + assert_eq!(series.len(), 6); + assert_eq!( + series + .iter() + .filter(|(_, _, value)| *value == 1) + .collect::>(), + vec![&("returning".to_string(), "primary".to_string(), 1)] + ); + assert_eq!(series.iter().map(|(_, _, value)| value).sum::(), 1); + } + + #[test] + fn returning_session_on_other_cluster_counts_fallback() { + let metrics = StargateMetrics::new().unwrap(); + let config = pulsar_wait_and_widen_config(None); + let hint = classified_hint(Some("B"), SessionState::Returning); + record_session_selection( + &metrics, + &config, + "pulsar-wait-and-widen", + "req-1", + &target(), + &hint, + "A", + 1, + ); + + assert!(session_series(&metrics).contains(&( + "returning".to_string(), + "fallback".to_string(), + 1 + ))); + } + + #[test] + fn new_session_primary_means_affinity_group() { + let config = pulsar_wait_and_widen_config(Some(2)); + for (state, rank_depth, expected) in [ + (SessionState::New, 1, "primary"), + (SessionState::New, 2, "primary"), + (SessionState::New, 3, "fallback"), + (SessionState::Stale, 3, "fallback"), + ] { + let metrics = StargateMetrics::new().unwrap(); + let hint = classified_hint(None, state); + record_session_selection( + &metrics, + &config, + "pulsar-wait-and-widen", + "req-1", + &target(), + &hint, + "A", + rank_depth, + ); + assert!( + session_series(&metrics).contains(&( + state.as_str().to_string(), + expected.to_string(), + 1 + )), + "{state:?} at rank {rank_depth}" + ); + } + } + + #[test] + fn unclassified_requests_record_nothing() { + let metrics = StargateMetrics::new().unwrap(); + let config = pulsar_wait_and_widen_config(None); + record_session_selection( + &metrics, + &config, + "pulsar-wait-and-widen", + "req-1", + &target(), + &hint("B"), + "B", + 1, + ); + + assert!(session_series(&metrics).is_empty()); + } +} diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs index de7bac7f43..de1ae4b089 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs @@ -13,11 +13,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::borrow::Cow; use std::collections::{HashMap, HashSet}; use std::time::{Duration, Instant}; use stargate_proto::pb::{InferenceServerStatus, ModelStats}; +use super::session::{self, MissingLastCluster}; use super::*; use crate::load_balancer::algorithm::{MAX_CACHE_AFFINITY_CACHE_KEY_BYTES, input_work_seconds}; use crate::load_balancer::pulsar::{PulsarLoadBalancer, pulsar_hash64, pulsar_ranked_indices}; @@ -91,6 +93,7 @@ fn request_with_priority<'a>( received_at: Instant::now(), request_slo: None, excluded_cluster_ids: None, + last_cluster: None, } } @@ -197,7 +200,7 @@ fn candidate(id: &str, kv_cache_free_tokens: u64) -> RoutedClusterSnapshot { } } -fn candidates(ids: &[&str]) -> Vec { +pub(super) fn candidates(ids: &[&str]) -> Vec { ids.iter().map(|id| candidate(id, 1024)).collect() } @@ -793,6 +796,42 @@ fn algorithm_specific_load_balancer_fields_are_rejected_for_other_algorithms() { } } +#[test] +fn last_cluster_affinity_is_rejected_for_other_algorithms() { + for algorithm in ["power-of-n", "round-robin", "random", "pulsar"] { + let raw = format!(r#"{{"algorithm":"{algorithm}","last_cluster_affinity":true}}"#); + let error = serde_json::from_str::(&raw) + .expect_err("last_cluster_affinity should be rejected"); + let message = error.to_string(); + assert!( + message.contains("last_cluster_affinity") && message.contains(algorithm), + "error should name the field and {algorithm}: {message}" + ); + } + for raw in [ + r#"{"default":"power-of-n","models":{"model-a":{"algorithm":"pulsar","last_cluster_affinity":true}}}"#, + r#"{"default":"power-of-n","request_algorithms":{"random":{"algorithm":"random","last_cluster_affinity":true}}}"#, + ] { + assert_json_rejected::(raw, "last_cluster_affinity"); + } +} + +#[test] +fn last_cluster_affinity_parses_for_wait_and_widen_algorithms() { + for algorithm in ["wait-and-widen", "pulsar-wait-and-widen"] { + let config: LoadBalancerAlgorithmConfig = parse_json(&format!( + r#"{{"algorithm":"{algorithm}","last_cluster_affinity":true}}"# + )); + assert_eq!( + config + .wait_and_widen_settings() + .and_then(|settings| settings.last_cluster_affinity), + Some(true) + ); + create_load_balancer_with_config(&config).expect("factory should accept the flag"); + } +} + #[test] fn power_of_n_sample_count_defaults_to_two() { let config = LoadBalancerAlgorithmConfig::from(LoadBalancerAlgorithm::PowerOfN); @@ -1436,6 +1475,7 @@ fn wait_and_widen_config_resolves_internal_defaults() { assert_eq!(config.max_queued, 0); assert!(config.ignore_queue_time); assert!(config.ignore_input_processing_time); + assert!(!config.last_cluster_affinity); let target = target(); let request = request(&target, None, Some(0)); @@ -3512,3 +3552,654 @@ fn wait_and_widen_without_affinity_prefers_a_fast_bucket_with_capacity() { ); } } + +// Session-aware wait-and-widen (`last_cluster_affinity`). + +const SESSION_WAIT: Duration = Duration::from_millis(300); + +fn session_wait_and_widen_config( + selection_count: usize, + last_cluster_affinity: bool, + configure: impl FnOnce(&mut WaitAndWidenAlgorithmConfig), +) -> WaitAndWidenConfig { + let mut config = wait_and_widen_affinity_algorithm_config(8, selection_count, Some(2)); + let settings = config.wait_and_widen_settings_mut().unwrap(); + settings.cache_affinity_wait_ms = Some(300); + settings.cache_affinity_input_tokens_scale = Some(0.1); + settings.last_cluster_affinity = Some(last_cluster_affinity); + configure(settings); + wait_and_widen_config(&config) +} + +fn session_request<'a>( + target: &'a RoutingTargetKey, + input_tokens: u64, + hint: &'a LastClusterHint, +) -> LoadBalancerRequest<'a> { + LoadBalancerRequest { + last_cluster: Some(hint), + ..request(target, Some("session-prefix"), Some(input_tokens)) + } +} + +fn last_cluster(id: &str) -> LastClusterHint { + LastClusterHint::new(Some(id.to_string())) +} + +fn mark_full(candidate: &mut RoutedClusterSnapshot) { + candidate.stats.max_engine_concurrency = 1; + candidate.stats.num_running_queries = 1; +} + +fn mark_free(candidate: &mut RoutedClusterSnapshot) { + candidate.stats.max_engine_concurrency = 1; + candidate.stats.num_running_queries = 0; +} + +fn selected_index(decision: LoadBalancerDecision) -> usize { + match decision { + LoadBalancerDecision::Selected(choice) => choice.candidate_index, + other => panic!("expected a selection, got {other:?}"), + } +} + +/// Affinity group of the session key, then every other candidate. +fn session_affinity_order( + config: &WaitAndWidenConfig, + target: &RoutingTargetKey, + candidates: &[RoutedClusterSnapshot], +) -> Vec { + let request = request(target, Some("session-prefix"), Some(0)); + let mut order = cache_affinity_candidate_indices(config, &request, candidates).unwrap(); + let rest: Vec<_> = (0..candidates.len()) + .filter(|index| !order.contains(index)) + .collect(); + order.extend(rest); + order +} + +#[test] +fn wait_and_widen_new_session_overflows_without_waiting() { + let config = session_wait_and_widen_config(1, true, |_| {}); + let target = target(); + let mut candidates = candidates(&["a", "b"]); + let order = session_affinity_order(&config, &target, &candidates); + mark_full(&mut candidates[order[0]]); + mark_free(&mut candidates[order[1]]); + + for hint in [LastClusterHint::default(), last_cluster("gone")] { + let lb = WaitAndWidenLoadBalancer::new(config.clone()); + let request = session_request(&target, 100, &hint); + assert_eq!( + selected_index(lb.decide_at(&request, &candidates, Duration::ZERO)), + order[1] + ); + } + + let flag_off = WaitAndWidenLoadBalancer::new(session_wait_and_widen_config(1, false, |_| {})); + let hint = LastClusterHint::default(); + assert_eq!( + flag_off.decide_at( + &session_request(&target, 100, &hint), + &candidates, + Duration::ZERO + ), + LoadBalancerDecision::Wait(SESSION_WAIT) + ); +} + +#[test] +fn wait_and_widen_new_session_matches_flag_off_without_wait_or_discount() { + let target = target(); + // Sampling covers every candidate and RTTs sit in distinct TTFT buckets, + // so each decision has one deterministic winner. + let config = session_wait_and_widen_config(2, true, |s| s.n = Some(4)); + let baseline = WaitAndWidenLoadBalancer::new(session_wait_and_widen_config(2, false, |s| { + s.n = Some(4); + s.cache_affinity_wait_ms = Some(0); + s.cache_affinity_input_tokens_scale = Some(1.0); + })); + let mut all_free = candidates(&["a", "b", "c", "d"]); + let order = session_affinity_order(&config, &target, &all_free); + for (rtt_ms, index) in [ + (120, order[0]), + (80, order[1]), + (40, order[2]), + (160, order[3]), + ] { + all_free[index].rtt = Duration::from_millis(rtt_ms); + } + let mut all_full = all_free.clone(); + all_full.iter_mut().for_each(mark_full); + let mut group_full = all_free.clone(); + mark_full(&mut group_full[order[0]]); + mark_full(&mut group_full[order[1]]); + + for candidates in [&all_free, &all_full, &group_full] { + for elapsed_ms in [0, 50, 300, 1_000] { + let elapsed = Duration::from_millis(elapsed_ms); + for hint in [LastClusterHint::default(), last_cluster("gone")] { + let lb = WaitAndWidenLoadBalancer::new(config.clone()); + let request = session_request(&target, 100, &hint); + assert_eq!( + lb.decide_at(&request, candidates, elapsed), + baseline.decide_at(&request, candidates, elapsed), + "at {elapsed_ms} ms" + ); + } + } + } + // The affinity group wins when it can serve; otherwise C overflows at once. + let hint = LastClusterHint::default(); + let request = session_request(&target, 100, &hint); + let lb = WaitAndWidenLoadBalancer::new(config); + assert_eq!( + selected_index(lb.decide_at(&request, &all_free, Duration::ZERO)), + order[1] + ); + assert_eq!( + selected_index(lb.decide_at(&request, &group_full, Duration::ZERO)), + order[2] + ); +} + +#[test] +fn wait_and_widen_new_session_uses_full_prefill_in_affinity_group() { + let target = target(); + let config = session_wait_and_widen_config(2, true, |_| {}); + let mut candidates = candidates(&["a", "b", "c", "d"]); + let order = session_affinity_order(&config, &target, &candidates); + // Full prefill favors the faster processor (1500 vs 2001 ms). + // Discounted prefill favors the closer backend (600 vs 201 ms). + candidates[order[0]].rtt = Duration::from_millis(500); + candidates[order[0]].stats.last_mean_input_tps = 1_000.0; + candidates[order[1]].rtt = Duration::from_millis(1); + candidates[order[1]].stats.last_mean_input_tps = 500.0; + + let flag_off = WaitAndWidenLoadBalancer::new(session_wait_and_widen_config(2, false, |_| {})); + let new_hint = LastClusterHint::default(); + let request = session_request(&target, 1_000, &new_hint); + assert_eq!( + selected_index(flag_off.decide_at(&request, &candidates, Duration::ZERO)), + order[1] + ); + for hint in [LastClusterHint::default(), last_cluster("gone")] { + let lb = WaitAndWidenLoadBalancer::new(config.clone()); + let request = session_request(&target, 1_000, &hint); + assert_eq!( + selected_index(lb.decide_at(&request, &candidates, Duration::ZERO)), + order[0] + ); + } +} + +#[test] +fn wait_and_widen_returning_session_waits_for_last_cluster() { + let target = target(); + let config = session_wait_and_widen_config(1, true, |_| {}); + let mut candidates = candidates(&["a", "b", "c"]); + let order = session_affinity_order(&config, &target, &candidates); + let (a, b) = (order[0], order[1]); + mark_free(&mut candidates[a]); + mark_full(&mut candidates[b]); + let hint = last_cluster(&candidates[b].cluster_id); + let lb = WaitAndWidenLoadBalancer::new(config); + let request = session_request(&target, 100, &hint); + + assert_eq!( + lb.decide_at(&request, &candidates, Duration::ZERO), + LoadBalancerDecision::Wait(SESSION_WAIT) + ); + for elapsed_ms in (0..300).step_by(10) { + assert!( + matches!( + lb.decide_at(&request, &candidates, Duration::from_millis(elapsed_ms)), + LoadBalancerDecision::Wait(_) + ), + "returning session must wait for its last cluster at {elapsed_ms} ms" + ); + } + assert_eq!( + hint.classification(), + Some(SessionClassification { + state: SessionState::Returning, + promoted: true, + }) + ); + + mark_free(&mut candidates[b]); + let choice = lb + .decide_at(&request, &candidates, Duration::ZERO) + .selected() + .unwrap(); + assert_eq!(choice.candidate_index, b); + assert_eq!(choice.rank_depth, 1); +} + +#[test] +fn wait_and_widen_returning_session_discounts_only_the_last_cluster() { + let target = target(); + let config = session_wait_and_widen_config(1, true, |_| {}); + let mut candidates = candidates(&["a", "b", "c"]); + let order = session_affinity_order(&config, &target, &candidates); + let (a, b, c) = (order[0], order[1], order[2]); + mark_full(&mut candidates[b]); + // A wins only when discounted (201 ms vs C at 1500 ms). Full prefill + // makes A 2001 ms, so C wins the open set. + candidates[a].rtt = Duration::from_millis(1); + candidates[a].stats.last_mean_input_tps = 500.0; + candidates[c].rtt = Duration::from_millis(500); + candidates[c].stats.last_mean_input_tps = 1_000.0; + let hint = last_cluster(&candidates[b].cluster_id); + let lb = WaitAndWidenLoadBalancer::new(config.clone()); + let request = session_request(&target, 1_000, &hint); + + assert_eq!( + selected_index(lb.decide_at(&request, &candidates, SESSION_WAIT)), + c + ); + + // Flag off, A is the discounted affinity member and wins at once. + let flag_off = WaitAndWidenLoadBalancer::new(session_wait_and_widen_config(1, false, |_| {})); + assert_eq!( + selected_index(flag_off.decide_at(&request, &candidates, Duration::ZERO)), + a + ); +} + +#[test] +fn wait_and_widen_returning_session_outside_group_replaces_last_member() { + let target = target(); + let config = session_wait_and_widen_config(2, true, |_| {}); + let mut candidates = candidates(&["a", "b", "c", "d"]); + let order = session_affinity_order(&config, &target, &candidates); + let (a, b, c, d) = (order[0], order[1], order[2], order[3]); + let hint = last_cluster(&candidates[c].cluster_id); + let lb = WaitAndWidenLoadBalancer::new(config); + let request = session_request(&target, 100, &hint); + + // The group is (C, A): B is not selectable during the wait. + for _ in 0..32 { + let choice = selected_index(lb.decide_at(&request, &candidates, Duration::ZERO)); + assert!([a, c].contains(&choice), "chose {choice}"); + } + + mark_full(&mut candidates[a]); + mark_full(&mut candidates[c]); + mark_full(&mut candidates[d]); + mark_free(&mut candidates[b]); + assert_eq!( + lb.decide_at(&request, &candidates, Duration::from_millis(100)), + LoadBalancerDecision::Wait(Duration::from_millis(200)) + ); + assert_eq!( + selected_index(lb.decide_at(&request, &candidates, SESSION_WAIT)), + b + ); +} + +#[test] +fn wait_and_widen_returning_session_inside_group_keeps_members() { + let target = target(); + let config = session_wait_and_widen_config(2, true, |_| {}); + let flag_off = WaitAndWidenLoadBalancer::new(session_wait_and_widen_config(2, false, |_| {})); + let mut candidates = candidates(&["a", "b", "c", "d"]); + let order = session_affinity_order(&config, &target, &candidates); + let hint = last_cluster(&candidates[order[1]].cluster_id); + let lb = WaitAndWidenLoadBalancer::new(config); + let request = session_request(&target, 100, &hint); + + for _ in 0..32 { + let choice = selected_index(lb.decide_at(&request, &candidates, Duration::ZERO)); + assert!(order[..2].contains(&choice), "chose {choice}"); + } + mark_full(&mut candidates[order[0]]); + mark_full(&mut candidates[order[1]]); + for elapsed_ms in [0, 100, 300] { + let elapsed = Duration::from_millis(elapsed_ms); + let session = lb.decide_at(&request, &candidates, elapsed); + let expected = flag_off.decide_at(&request, &candidates, elapsed); + assert_eq!( + matches!(session, LoadBalancerDecision::Wait(_)), + matches!(expected, LoadBalancerDecision::Wait(_)), + "{session:?} vs {expected:?} at {elapsed_ms} ms" + ); + } +} + +#[test] +fn wait_and_widen_retry_excluded_last_cluster_acts_as_ineligible_primary() { + let target = target(); + let config = session_wait_and_widen_config(1, true, |_| {}); + let candidates = candidates(&["a", "b", "c"]); + let order = session_affinity_order(&config, &target, &candidates); + let b_id = candidates[order[1]].cluster_id.clone(); + let hint = last_cluster(&b_id); + let excluded = excluded(&[&b_id]); + let lb = WaitAndWidenLoadBalancer::new(config); + let request = LoadBalancerRequest { + excluded_cluster_ids: Some(&excluded), + ..session_request(&target, 100, &hint) + }; + + assert_eq!( + lb.decide_at(&request, &candidates, Duration::ZERO), + LoadBalancerDecision::Wait(SESSION_WAIT) + ); + let choice = selected_index(lb.decide_at(&request, &candidates, SESSION_WAIT)); + assert_ne!(candidates[choice].cluster_id, b_id); +} + +#[test] +fn wait_and_widen_promotion_is_not_cached() { + let target = target(); + let config = session_wait_and_widen_config(1, true, |_| {}); + let candidates = candidates(&["a", "b", "c"]); + let order = session_affinity_order(&config, &target, &candidates); + let lb = WaitAndWidenLoadBalancer::new(config.clone()); + let returning = last_cluster(&candidates[order[1]].cluster_id); + assert_eq!( + selected_index(lb.decide_at( + &session_request(&target, 100, &returning), + &candidates, + Duration::ZERO + )), + order[1] + ); + + let new_session = LastClusterHint::default(); + assert_eq!( + selected_index(lb.decide_at( + &session_request(&target, 100, &new_session), + &candidates, + Duration::ZERO + )), + order[0] + ); + assert_eq!( + cache_affinity_candidate_indices( + &config, + &request(&target, Some("session-prefix"), Some(0)), + &candidates + ) + .unwrap(), + vec![order[0]] + ); +} + +#[test] +fn wait_and_widen_session_decisions_match_across_instances() { + let target = target(); + let config = session_wait_and_widen_config(2, true, |s| s.n = Some(4)); + let mut candidates = candidates(&["a", "b", "c", "d"]); + let order = session_affinity_order(&config, &target, &candidates); + let (a, b, c, d) = (order[0], order[1], order[2], order[3]); + // B is fastest but leaves the group when C is promoted: (C, A). + for (rtt_ms, index) in [(60, a), (20, b), (100, c), (140, d)] { + candidates[index].rtt = Duration::from_millis(rtt_ms); + } + let last_id = candidates[c].cluster_id.clone(); + let mut group_full = candidates.clone(); + mark_full(&mut group_full[a]); + mark_full(&mut group_full[c]); + + let instances: Vec<_> = (0..2) + .map(|_| WaitAndWidenLoadBalancer::new(config.clone())) + .collect(); + let plans: Vec<_> = instances + .iter() + .map(|_| { + let hint = last_cluster(&last_id); + let request = session_request(&target, 100, &hint); + let affinity = + cache_affinity_candidate_indices(&config, &request, &candidates).unwrap(); + let plan = session::plan_affinity( + config.last_cluster_affinity, + &request, + &candidates, + Cow::Borrowed(affinity.as_slice()), + MissingLastCluster::InsertAndTruncate, + config.cache_affinity_wait, + config.cache_affinity_input_tokens_scale, + ); + ( + plan.order.into_owned(), + plan.affinity_wait, + plan.input_tokens_scale, + hint.classification(), + ) + }) + .collect(); + assert_eq!(plans[0], plans[1]); + assert_eq!(plans[0].0, vec![c, a]); + assert_eq!((plans[0].1, plans[0].2), (SESSION_WAIT, 0.1)); + + for (snapshot, elapsed_ms, expected) in [ + (&candidates, 0, Some(a)), + (&group_full, 150, None), + (&group_full, 300, Some(b)), + (&group_full, 1_000, Some(b)), + ] { + let elapsed = Duration::from_millis(elapsed_ms); + let decisions: Vec<_> = instances + .iter() + .map(|lb| { + let hint = last_cluster(&last_id); + lb.decide_at(&session_request(&target, 100, &hint), snapshot, elapsed) + }) + .collect(); + assert_eq!(decisions[0], decisions[1], "at {elapsed_ms} ms"); + assert_eq!( + decisions[0].selected().map(|choice| choice.candidate_index), + expected, + "at {elapsed_ms} ms" + ); + } +} + +#[test] +fn wait_and_widen_ignores_last_cluster_when_flag_is_off() { + let target = target(); + let config = session_wait_and_widen_config(1, false, |_| {}); + let mut candidates = candidates(&["a", "b", "c"]); + let order = session_affinity_order(&config, &target, &candidates); + mark_full(&mut candidates[order[1]]); + let lb = WaitAndWidenLoadBalancer::new(config); + + for elapsed_ms in [0, 300] { + let elapsed = Duration::from_millis(elapsed_ms); + let hint = last_cluster(&candidates[order[1]].cluster_id); + let with_header = lb.decide_at(&session_request(&target, 100, &hint), &candidates, elapsed); + let without = lb.decide_at( + &request(&target, Some("session-prefix"), Some(100)), + &candidates, + elapsed, + ); + assert_eq!(with_header, without); + assert_eq!(hint.classification(), None); + } +} + +#[test] +fn wait_and_widen_session_needs_affinity_key_and_group() { + let target = target(); + let candidates = candidates(&["a", "b"]); + let hint = last_cluster("b"); + + let without_key = WaitAndWidenLoadBalancer::new(session_wait_and_widen_config(1, true, |_| {})); + let request_without_key = LoadBalancerRequest { + last_cluster: Some(&hint), + ..request(&target, None, Some(100)) + }; + assert!( + without_key + .decide_at(&request_without_key, &candidates, Duration::ZERO) + .selected() + .is_some() + ); + assert_eq!(hint.classification(), None); + + let without_group = + WaitAndWidenLoadBalancer::new(session_wait_and_widen_config(1, true, |s| { + s.cache_affinity_backend_selection_count = None; + })); + assert!( + without_group + .decide_at( + &session_request(&target, 100, &hint), + &candidates, + Duration::ZERO + ) + .selected() + .is_some() + ); + assert_eq!(hint.classification(), None); +} + +#[test] +fn router_ignores_last_cluster_for_algorithm_without_flag() { + let router = router_from_json( + r#"{ + "default": "power-of-n", + "models": { + "model-a": { + "algorithm": "wait-and-widen", + "seed": "seed-1", + "cache_affinity_backend_selection_count": 1, + "cache_affinity_wait_ms": 300, + "request_algorithms": { + "pulsar-wait-and-widen": { + "algorithm": "pulsar-wait-and-widen", + "last_cluster_affinity": true + } + } + } + } + }"#, + ); + let resolution = router + .resolve_algorithm_override("model-a", None) + .expect("default model algorithm should resolve"); + assert_eq!( + resolution + .config() + .wait_and_widen_settings() + .and_then(|settings| settings.last_cluster_affinity), + None + ); + let resolution = router + .resolve_algorithm_override( + "model-a", + Some( + &"pulsar-wait-and-widen" + .parse::() + .unwrap(), + ), + ) + .expect("override should resolve"); + assert_eq!( + resolution + .config() + .wait_and_widen_settings() + .and_then(|settings| settings.last_cluster_affinity), + Some(true) + ); +} + +#[test] +fn routing_method_override_without_flag_ignores_last_cluster() { + let router = router_from_json( + r#"{ + "default": "power-of-n", + "models": { + "model-a": { + "algorithm": "wait-and-widen", + "seed": "seed-1", + "cache_affinity_virtual_nodes": 8, + "cache_affinity_backend_selection_count": 1, + "cache_affinity_wait_ms": 300, + "last_cluster_affinity": true, + "request_algorithms": { + "pulsar-wait-and-widen": { + "algorithm": "pulsar-wait-and-widen", + "seed": "seed-1", + "cache_affinity_wait_ms": 300 + } + } + } + } + }"#, + ); + let algorithm_override = "pulsar-wait-and-widen" + .parse::() + .unwrap(); + let resolution = router + .resolve_algorithm_override("model-a", Some(&algorithm_override)) + .expect("override should resolve"); + let target = target(); + let target_state = LoadBalancerTargetState::default(); + let metrics = crate::metrics::StargateMetrics::new().unwrap(); + let base = candidates(&["a", "b", "c"]); + let mut snapshots = vec![base.clone()]; + for index in 0..base.len() { + let mut snapshot = base.clone(); + mark_full(&mut snapshot[index]); + snapshots.push(snapshot); + } + + for snapshot in &snapshots { + for last in &base { + let hint = last_cluster(&last.cluster_id); + let with_header = router.decide_with_algorithm_resolution( + &target_state, + &session_request(&target, 100, &hint), + snapshot, + &resolution, + ); + let without = router.decide_with_algorithm_resolution( + &target_state, + &request(&target, Some("session-prefix"), Some(100)), + snapshot, + &resolution, + ); + // Decisions use wall-clock elapsed time, so compare wait kinds. + let comparable = |decision: LoadBalancerDecision| match decision { + LoadBalancerDecision::Wait(_) => LoadBalancerDecision::Wait(Duration::ZERO), + decision => decision, + }; + assert_eq!(comparable(with_header), comparable(without)); + assert_eq!(hint.classification(), None); + record_session_selection( + &metrics, + resolution.config(), + "pulsar-wait-and-widen", + "req-1", + &target, + &hint, + &last.cluster_id, + 1, + ); + } + } + assert!( + metrics + .registry() + .gather() + .iter() + .all(|family| family.name() != "stargate_routing_session_selections_total") + ); + + // The model default carries the flag and classifies the same request. + let default_resolution = router.resolve_algorithm_override("model-a", None).unwrap(); + let hint = last_cluster("b"); + let _ = router.decide_with_algorithm_resolution( + &target_state, + &session_request(&target, 100, &hint), + &base, + &default_resolution, + ); + assert_eq!( + hint.classification().map(|c| c.state), + Some(SessionState::Returning) + ); +} diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/wait_and_widen.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/wait_and_widen.rs index 1f139b7767..80691e7f94 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/wait_and_widen.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/wait_and_widen.rs @@ -17,12 +17,14 @@ mod cache_affinity; mod estimates; mod fast_path; +use std::borrow::Cow; use std::cmp::Ordering; use std::ops::RangeInclusive; use std::time::Duration; use rand::Rng; +use super::session::{self, MissingLastCluster}; use super::{ ClusterComparator, LoadBalancer, LoadBalancerAlgorithmConfig, LoadBalancerCandidateChoice, LoadBalancerDecision, LoadBalancerRequest, Ttft, ttft, @@ -65,6 +67,7 @@ pub(super) struct WaitAndWidenConfig { pub(super) max_queued: u64, pub(super) ignore_queue_time: bool, pub(super) ignore_input_processing_time: bool, + pub(super) last_cluster_affinity: bool, } impl WaitAndWidenConfig { @@ -100,6 +103,7 @@ impl WaitAndWidenConfig { max_queued: config.max_queued.unwrap_or(0), ignore_queue_time: config.ignore_queue_time.unwrap_or(false), ignore_input_processing_time: config.ignore_input_processing_time.unwrap_or(false), + last_cluster_affinity: config.last_cluster_affinity.unwrap_or(false), }) } @@ -179,14 +183,25 @@ impl WaitAndWidenLoadBalancer { self.cache_affinity .candidate_indices(&self.config, request, candidates) { + // Session rules adjust a per-request copy; the ring cache keeps + // the original selection. + let plan = session::plan_affinity( + self.config.last_cluster_affinity, + request, + candidates, + Cow::Borrowed(affinity_indices.as_slice()), + MissingLastCluster::InsertAndTruncate, + self.config.cache_affinity_wait, + self.config.cache_affinity_input_tokens_scale, + ); // Every attempt checks affinity first, even after the hold expires. // Affinity uses the same TTFT buckets, admission and comparator as // global routing. A busy ring primary must not bypass those rules. let affinity_bucket_wait = match self.decide_from_candidate_indices( request, candidates, - &affinity_indices, - self.config.cache_affinity_input_tokens_scale, + &plan.order, + plan.input_tokens_scale, elapsed, ) { LoadBalancerDecision::Selected(choice) => { @@ -195,8 +210,8 @@ impl WaitAndWidenLoadBalancer { LoadBalancerDecision::Wait(delay) => Some(delay), LoadBalancerDecision::Unavailable => None, }; - if elapsed < self.config.cache_affinity_wait { - let remaining = self.config.cache_affinity_wait - elapsed; + if elapsed < plan.affinity_wait { + let remaining = plan.affinity_wait - elapsed; return LoadBalancerDecision::Wait( affinity_bucket_wait.map_or(remaining, |delay| delay.min(remaining)), ); @@ -209,10 +224,10 @@ impl WaitAndWidenLoadBalancer { candidates.iter(), candidates, 1.0, - elapsed.saturating_sub(self.config.cache_affinity_wait), + elapsed.saturating_sub(plan.affinity_wait), ) { LoadBalancerDecision::Selected(mut choice) => { - choice.rank_depth = affinity_indices.len() + 1; + choice.rank_depth = plan.order.len() + 1; LoadBalancerDecision::Selected(choice) } LoadBalancerDecision::Wait(delay) => LoadBalancerDecision::Wait( @@ -246,10 +261,6 @@ impl WaitAndWidenLoadBalancer { } } - pub(super) fn has_queue_slo(&self, request: &LoadBalancerRequest<'_>) -> bool { - self.config.max_queue_time(request).is_some() - } - #[cfg(test)] pub(super) fn cached_affinity_key_bytes(&self) -> usize { self.cache_affinity.cached_key_bytes() diff --git a/src/libraries/rust/stargate/crates/stargate/src/metrics.rs b/src/libraries/rust/stargate/crates/stargate/src/metrics.rs index e425c4813f..31d289f051 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/metrics.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/metrics.rs @@ -102,6 +102,7 @@ define_stargate_metrics! { proxy_retries_total("proxy_retries_total", "Total number of proxy retries", ["routing_key", "model", "reason"]); routing_selections_total("routing_selections_total", "Total number of primary and ranked fallback cluster choices used for upstream attempts", ["routing_key", "model", "algorithm", "selection"]); routing_kv_free_token_fallback_selections_total("routing_kv_free_token_fallback_selections_total", "Total number of selected routes reached after a higher-ranked candidate was skipped by KV free-token eligibility", ["routing_key", "model", "algorithm"]); + routing_session_selections_total("routing_session_selections_total", "Total number of cluster choices for session-classified requests by session state and whether the session's primary target was chosen", ["routing_key", "model", "algorithm", "session_state", "selection"]); proxy_retry_exhausted_total("proxy_retry_exhausted_total", "Total number of proxy requests that exhausted retry options", ["routing_key", "model", "reason"]); proxy_ambiguous_delivery_total("proxy_ambiguous_delivery_total", "Total number of proxy requests stopped because delivery may have occurred", []); admission_rejections_total("admission_rejections_total", "Total number of requests rejected by local admission control", ["routing_key", "model", "reason"]); @@ -206,6 +207,7 @@ impl StargateMetrics { GenericCounter, proxy_retries_total(routing_key: Option<&str>, model: &str, reason: &str) => [routing_key.unwrap_or(""), model, reason]; GenericCounter, routing_selections_total(routing_key: Option<&str>, model: &str, algorithm: &str, selection: &str) => [routing_key.unwrap_or(""), model, algorithm, selection]; GenericCounter, routing_kv_free_token_fallback_selections_total(routing_key: Option<&str>, model: &str, algorithm: &str) => [routing_key.unwrap_or(""), model, algorithm]; + GenericCounter, routing_session_selections_total(routing_key: Option<&str>, model: &str, algorithm: &str, session_state: &str, selection: &str) => [routing_key.unwrap_or(""), model, algorithm, session_state, selection]; GenericCounter, proxy_retry_exhausted_total(routing_key: Option<&str>, model: &str, reason: &str) => [routing_key.unwrap_or(""), model, reason]; GenericCounter, proxy_ambiguous_delivery_total() => []; GenericCounter, admission_rejections_total(routing_key: Option<&str>, model: &str, reason: &str) => [routing_key.unwrap_or(""), model, reason]; diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs index 7f8f191e14..10dbb98987 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs @@ -2136,6 +2136,7 @@ async fn routing_target_recreation_starts_fresh_load_balancer_state() { received_at: Instant::now(), request_slo: None, excluded_cluster_ids: None, + last_cluster: None, }; let candidates = [ RoutedClusterSnapshot { @@ -2364,6 +2365,7 @@ async fn three_shared_backends_merge_sparse_priority_maps() { received_at: Instant::now(), request_slo: None, excluded_cluster_ids: None, + last_cluster: None, }; let choice = router .choose_candidate(snapshot.load_balancers(), &request, snapshot.clusters()) @@ -2621,6 +2623,7 @@ async fn routing_uses_load_from_every_shared_cluster_backend() { received_at: Instant::now(), request_slo: None, excluded_cluster_ids: None, + last_cluster: None, }; let choice = router .choose_candidate(snapshot.load_balancers(), &request, snapshot.clusters()) @@ -2726,6 +2729,7 @@ async fn registered_backend_rtt_means_drive_cluster_load_balancer_selection() { received_at: Instant::now(), request_slo: None, excluded_cluster_ids: None, + last_cluster: None, }; let choice = router .choose_candidate( diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs index fa9c1fa83d..9037ff85c2 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs @@ -1587,3 +1587,279 @@ async fn wait_and_widen_capacity_rejection_returns_overloaded_error() { stop_backends(&mut backends); stargate.shutdown().await; } + +const LAST_CLUSTER_HEADER: &str = "x-stargate-last-cluster-id"; + +fn session_selections(handle: &StargateHandle, model: &str, labels: &[(&str, &str)]) -> f64 { + handle + .metrics() + .registry() + .gather() + .iter() + .filter(|family| family.name() == "stargate_routing_session_selections_total") + .flat_map(|family| family.get_metric()) + .filter(|metric| { + let has_label = |name: &str, value: &str| { + metric + .get_label() + .iter() + .any(|label| label.name() == name && label.value() == value) + }; + has_label("model", model) && labels.iter().all(|(name, value)| has_label(name, value)) + }) + .map(|metric| metric.get_counter().value()) + .sum() +} + +fn session_stats(num_running_queries: u64, max_engine_concurrency: u64) -> CurrentModelStats { + CurrentModelStats { + last_mean_input_tps: 100.0, + max_input_tps: Some(100.0), + max_output_tps: 1000.0, + num_running_queries, + max_engine_concurrency: Some(max_engine_concurrency), + ..CurrentModelStats::default() + } +} + +async fn wait_for_clusters( + state: &StargateState, + model_id: &str, + description: &str, + ready: impl Fn(&[RoutedClusterSnapshot]) -> bool, +) { + let target = RoutingTargetKey { + routing_key: None, + model_id: model_id.to_string(), + }; + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + loop { + let clusters = state.cluster_candidates_for_target(&target).await; + if ready(&clusters) { + return; + } + assert!( + tokio::time::Instant::now() < deadline, + "model '{model_id}' clusters did not reach {description}: {clusters:?}" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +fn cluster_has_capacity(clusters: &[RoutedClusterSnapshot], cluster_id: &str, max: u64) -> bool { + clusters.iter().any(|cluster| { + cluster.cluster_id == cluster_id && cluster.stats.max_engine_concurrency == max + }) +} + +async fn send_session_request( + chat: &ChatRequests, + request_id: &str, + affinity_key: &str, + last_cluster: Option<&str>, +) -> String { + let mut request = chat.with_affinity(request_id, affinity_key); + if let Some(last_cluster) = last_cluster { + request = request.header(LAST_CLUSTER_HEADER, last_cluster); + } + let response = request.send().await.expect("request failed"); + assert_eq!(response.status(), 200, "request {request_id} failed"); + let cluster_id = response_header(&response, "x-stargate-cluster-id").to_string(); + let _ = tokio::time::timeout(Duration::from_secs(15), response.bytes()).await; + cluster_id +} + +struct CacheThrashRun { + model: String, + rank_one: String, + rank_two: String, + follow_ups: Vec, + stargate: RunningStargate, + backends: Vec, +} + +impl CacheThrashRun { + fn session_selections(&self, labels: &[(&str, &str)]) -> f64 { + session_selections(&self.stargate.handle, &self.model, labels) + } + + async fn finish(mut self) { + stop_backends(&mut self.backends); + self.stargate.shutdown().await; + } +} + +/// Reproduces the overflow cache-thrash scenario for one algorithm and flag +/// value. Records the rank-1 cluster, the overflow cluster, and the clusters +/// chosen for follow-up requests that carry the overflow cluster as the hint. +async fn run_cache_thrash(algorithm: &str, last_cluster_affinity: bool) -> CacheThrashRun { + let flag = if last_cluster_affinity { "on" } else { "off" }; + let model = format!("{algorithm}-thrash-{flag}"); + let virtual_nodes = if algorithm == "wait-and-widen" { + r#""cache_affinity_virtual_nodes": 8,"# + } else { + "" + }; + let config = format!( + r#"{{ + "default": "power-of-n", + "models": {{ + "{model}": {{ + "algorithm": "{algorithm}", + "seed": "test-seed", + "require_cache_affinity_key": true, + {virtual_nodes} + "cache_affinity_backend_selection_count": 1, + "cache_affinity_wait_ms": 300, + "cache_affinity_input_tokens_scale": 0.1, + "last_cluster_affinity": {last_cluster_affinity} + }} + }} + }}"# + ); + let stargate = RunningStargate::start(&format!("test-sg-{model}"), Some(&config)).await; + let cluster_ids = [format!("{model}-x"), format!("{model}-y")]; + let mut backends = Vec::new(); + for cluster_id in &cluster_ids { + let backend = + RegisteredBackend::active_with_fast_updates(stargate.grpc_addr, &model, cluster_id) + .await; + backend.set_stats(session_stats(0, 100)); + backends.push(backend); + } + let state = stargate.handle.state(); + wait_for_clusters(&state, &model, "two free clusters", |clusters| { + clusters.len() == 2 + && cluster_ids + .iter() + .all(|cluster_id| cluster_has_capacity(clusters, cluster_id, 100)) + }) + .await; + wait_for_routing_with_cache_affinity( + stargate.http_addr, + &model, + "warmup-session", + Duration::from_secs(5), + ) + .await; + + let chat = ChatRequests::new(stargate.http_addr, &model); + let affinity_key = "thrash-session"; + + // Rank 1 for the session is stable while every cluster is free. + let rank_one = send_session_request(&chat, "thrash-probe-0", affinity_key, None).await; + for attempt in 1..5 { + let chosen = send_session_request( + &chat, + &format!("thrash-probe-{attempt}"), + affinity_key, + None, + ) + .await; + assert_eq!(chosen, rank_one, "rank 1 must be stable for a free pool"); + } + let rank_one_index = cluster_ids + .iter() + .position(|cluster_id| *cluster_id == rank_one) + .expect("rank 1 must be a registered cluster"); + let rank_two = cluster_ids[1 - rank_one_index].clone(); + + // Fill rank 1. The first request of the session overflows to rank 2. + backends[rank_one_index].set_stats(session_stats(1, 1)); + wait_for_clusters(&state, &model, "full rank 1", |clusters| { + clusters.iter().any(|cluster| { + cluster.cluster_id == rank_one + && cluster.stats.max_engine_concurrency == 1 + && cluster.stats.num_running_queries >= 1 + }) + }) + .await; + let overflow = send_session_request(&chat, "thrash-first", affinity_key, None).await; + assert_eq!(overflow, rank_two, "a full rank 1 must overflow to rank 2"); + + // Rank 1 drains. Follow-up requests carry the overflow cluster as the hint. + backends[rank_one_index].set_stats(session_stats(0, 100)); + wait_for_clusters(&state, &model, "drained rank 1", |clusters| { + cluster_has_capacity(clusters, &rank_one, 100) + }) + .await; + let mut follow_ups = Vec::new(); + for attempt in 0..10 { + follow_ups.push( + send_session_request( + &chat, + &format!("thrash-follow-up-{attempt}"), + affinity_key, + Some(&overflow), + ) + .await, + ); + } + + CacheThrashRun { + model, + rank_one, + rank_two, + follow_ups, + stargate, + backends, + } +} + +async fn assert_cache_thrash_keeps_session_on_last_cluster(algorithm: &str) { + let run = run_cache_thrash(algorithm, true).await; + assert_eq!( + run.follow_ups, + vec![run.rank_two.clone(); 10], + "{algorithm}: returning requests must stay on their last cluster" + ); + assert_eq!( + run.session_selections(&[("session_state", "returning"), ("selection", "primary")]), + 10.0 + ); + assert_eq!( + run.session_selections(&[("session_state", "returning"), ("selection", "fallback")]), + 0.0 + ); + assert_eq!( + run.session_selections(&[("session_state", "new"), ("selection", "fallback")]), + 1.0, + "{algorithm}: only the overflowing first request leaves the affinity group" + ); + run.finish().await; +} + +async fn assert_cache_thrash_without_flag_returns_to_rank_one(algorithm: &str) { + let run = run_cache_thrash(algorithm, false).await; + assert_eq!( + run.follow_ups, + vec![run.rank_one.clone(); 10], + "{algorithm}: without last_cluster_affinity the hint is ignored" + ); + assert_eq!( + run.session_selections(&[]), + 0.0, + "{algorithm}: unclassified requests must not record session selections" + ); + run.finish().await; +} + +#[tokio::test] +async fn wait_and_widen_last_cluster_affinity_prevents_cache_thrash() { + assert_cache_thrash_keeps_session_on_last_cluster("wait-and-widen").await; +} + +#[tokio::test] +async fn wait_and_widen_without_last_cluster_affinity_returns_to_rank_one() { + assert_cache_thrash_without_flag_returns_to_rank_one("wait-and-widen").await; +} + +#[tokio::test] +async fn pulsar_wait_and_widen_last_cluster_affinity_prevents_cache_thrash() { + assert_cache_thrash_keeps_session_on_last_cluster("pulsar-wait-and-widen").await; +} + +#[tokio::test] +async fn pulsar_wait_and_widen_without_last_cluster_affinity_returns_to_rank_one() { + assert_cache_thrash_without_flag_returns_to_rank_one("pulsar-wait-and-widen").await; +} diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs index 1603f11e21..7f1eeb7944 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs @@ -3713,3 +3713,183 @@ async fn exercise_native_messages_contract(protocol: TunnelTransportProtocol, re fixture.shutdown().await; backend_task.abort(); } + +#[tokio::test] +async fn last_cluster_hint_is_consumed_and_classifies_only_flagged_affinity_requests() { + const LAST_CLUSTER_HEADER: &str = "x-stargate-last-cluster-id"; + const SESSION_METRIC: &str = "stargate_routing_session_selections_total"; + let flagged_model = "last-cluster-on-model"; + let unflagged_model = "last-cluster-off-model"; + let cluster_id = "last-cluster-cluster"; + + init_crypto(); + let mut config_file = tempfile::NamedTempFile::new().expect("failed to create temp file"); + std::io::Write::write_all( + &mut config_file, + br#"{ + "default": "power-of-n", + "request_algorithms": {"power-of-n": "power-of-n"}, + "models": { + "last-cluster-on-model": { + "algorithm": "wait-and-widen", + "seed": "test-seed", + "cache_affinity_backend_selection_count": 1, + "last_cluster_affinity": true + }, + "last-cluster-off-model": { + "algorithm": "wait-and-widen", + "seed": "test-seed", + "cache_affinity_backend_selection_count": 1 + } + } + }"#, + ) + .expect("failed to write config"); + let config_path = config_file.path().to_str().unwrap().to_string(); + let (grpc_addr, http_addr, runtime) = + make_stargate_runtime_with_lb("last-cluster-header", Some(config_path)); + let handle = runtime.start().await.expect("stargate failed to start"); + let mut fixture = ProxyFixture::new(grpc_addr, http_addr, handle); + + let captured = Arc::new(std::sync::Mutex::new(Vec::::new())); + let captured_for_app = captured.clone(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let upstream_addr = listener.local_addr().unwrap(); + let app = Router::new() + .route( + "/v1/chat/completions", + post(move |req: Request| { + let captured = captured_for_app.clone(); + async move { + captured + .lock() + .expect("header capture poisoned") + .push(req.headers().clone()); + Response::builder() + .status(StatusCode::OK) + .header("content-type", "text/event-stream") + .body(Body::from( + "data: {\"object\":\"chat.completion.chunk\"}\n\ndata: [DONE]\n\n", + )) + .expect("success response should build") + } + }), + ) + .route("/health", get(|| async { "ok" })); + let backend_task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let tunnel = start_quic_http_tunnel(QuicHttpTunnelConfig::new( + "127.0.0.1:0".parse().unwrap(), + format!("http://{upstream_addr}"), + )) + .await + .unwrap(); + let runtime_state = PylonRuntimeState::new( + InferenceServerStatus::Active, + &[flagged_model.to_string(), unflagged_model.to_string()], + ); + for model in [flagged_model, unflagged_model] { + runtime_state.set_model_stats( + model, + CurrentModelStats { + last_mean_input_tps: 1000.0, + max_input_tps: Some(1000.0), + ..CurrentModelStats::default() + }, + ); + } + fixture.register( + active_registration_config_with_state( + grpc_addr, + "last-cluster-backend", + cluster_id, + format!("quic://{}", tunnel.listen_addr()), + format!("http://{upstream_addr}"), + runtime_state, + ), + "registration failed", + ); + fixture.own_tunnel(tunnel); + fixture.wait_for_clusters(flagged_model, 1).await; + fixture.wait_for_clusters(unflagged_model, 1).await; + for model in [flagged_model, unflagged_model] { + wait_for_routing(http_addr, model, Duration::from_secs(5)).await; + } + + // (model, routing method override, cache affinity key, classified) + let cases = [ + (flagged_model, None, Some("session-a"), true), + (flagged_model, None, None, false), + (flagged_model, Some("power-of-n"), Some("session-a"), false), + (unflagged_model, None, Some("session-a"), false), + ]; + let client = reqwest::Client::new(); + for (index, (model, routing_method, affinity_key, classified)) in cases.into_iter().enumerate() + { + let context = format!( + "model={model} routing_method={routing_method:?} affinity_key={affinity_key:?}" + ); + let request_id = format!("last-cluster-request-{index}"); + let model_label = format!(r#"model="{model}""#); + let before = metrics_text(fixture.handle.metrics().registry()); + let mut request = proxy_json_request( + &client, + http_addr, + "/v1/chat/completions", + model, + &request_id, + &streaming_chat_body(model), + ) + .header(LAST_CLUSTER_HEADER, cluster_id); + if let Some(routing_method) = routing_method { + request = request.header("x-routing-method", routing_method); + } + if let Some(affinity_key) = affinity_key { + request = request.header("x-cache-affinity-key", affinity_key); + } + let response = request.send().await.expect("request failed"); + assert_eq!(response.status(), StatusCode::OK, "{context}"); + assert_eq!( + response_header(&response, "x-stargate-cluster-id"), + Some(cluster_id), + "{context}" + ); + let _ = response.bytes().await; + + let headers = captured + .lock() + .expect("header capture poisoned") + .iter() + .rev() + .find(|headers| { + headers + .get("x-request-id") + .is_some_and(|value| value.as_bytes() == request_id.as_bytes()) + }) + .cloned() + .unwrap_or_else(|| panic!("upstream did not receive {request_id}: {context}")); + assert!( + !headers.contains_key(LAST_CLUSTER_HEADER), + "{LAST_CLUSTER_HEADER} leaked upstream: {context}" + ); + + let after = metrics_text(fixture.handle.metrics().registry()); + let expected = if classified { 1.0 } else { 0.0 }; + assert_metric_delta(&before, &after, SESSION_METRIC, &[&model_label], expected); + assert_metric_delta( + &before, + &after, + SESSION_METRIC, + &[ + &model_label, + r#"session_state="returning""#, + r#"selection="primary""#, + ], + expected, + ); + } + + fixture.shutdown().await; + backend_task.abort(); +} diff --git a/src/libraries/rust/stargate/docs/load-balancer-configuration.md b/src/libraries/rust/stargate/docs/load-balancer-configuration.md index 6f1637701d..52ec9f3950 100644 --- a/src/libraries/rust/stargate/docs/load-balancer-configuration.md +++ b/src/libraries/rust/stargate/docs/load-balancer-configuration.md @@ -111,7 +111,7 @@ Choose based on the routing goal and available backend statistics: | Compare a small random sample using TTFT or another load signal. | `power-of-n` | Signals required by the configured comparator. | | Minimize estimated time to first token across heterogeneous or remote clusters while controlling which TTFT bands are eligible. | `wait-and-widen` | Forwarded health RTT and model statistics. Valid `last_mean_input_tps` is needed when queued or request input work is nonzero. | | Keep the same prefix on a stable, capacity-weighted cluster. | `pulsar` | Positive finite `max_input_tps` for every participating cluster by default. | -| Keep Pulsar affinity when possible, then widen through the Pulsar ranking when the primary cannot meet queue policy. | `pulsar-wait-and-widen` | Pulsar capacity plus the RTT and queue statistics used by `wait-and-widen`. | +| Keep Pulsar affinity when possible, then widen through the Pulsar ranking when the primary is full or cannot meet queue policy. | `pulsar-wait-and-widen` | Pulsar capacity plus the RTT and queue statistics used by `wait-and-widen`. | Use `round-robin` for deterministic cycling and `random` for uniform random selection when routing should not depend on backend load statistics. @@ -302,11 +302,12 @@ Minimal configuration: ## `pulsar-wait-and-widen` `pulsar-wait-and-widen` is `wait-and-widen` with the Pulsar ranking as the -source of affinity. Without queue-SLO fields, an eligible Pulsar primary wins -immediately, so the affinity wait, band widening, and `fallback_max_queued` -below apply only when queue-SLO fields are set or the primary is ineligible. +source of affinity. Every request goes through the affinity phase described +below, including a request whose rank-1 cluster is eligible: queue admission +with `max_queued`, the affinity wait, and the affinity prefill discount. The +rank-1 cluster receives no request without a load check. -In that case the affinity group is the top +The affinity group is the top `cache_affinity_backend_selection_count` clusters of the Pulsar ranking. The count defaults to `1`, the primary, and `0` also means `1`. WaitAndWiden selection runs within the group, and `cache_affinity_input_tokens_scale` @@ -359,6 +360,147 @@ Minimal configuration: } ``` +### Upgrade from the eligible-primary shortcut + +Earlier releases sent a request straight to an eligible rank-1 cluster when +the configuration set neither `max_queue_time_floor_ms` nor +`max_queue_time_ceil_ms`. That path ignored load, so a full rank-1 cluster kept +taking requests without limit, and a hot affinity key slowed every other key +that shared its rank-1 cluster. + +That shortcut is removed. A configuration without queue-SLO fields that keeps +the defaults (`max_queued: 0`, `cache_affinity_wait_ms: 0`) now overflows as +soon as its rank-1 cluster has no free engine slot. Before upgrading, review +`pulsar-wait-and-widen` configurations without queue-SLO fields: + +- For bounded stickiness, raise `max_queued` for the affinity group and set + `cache_affinity_wait_ms`. +- For unlimited stickiness, use `pulsar`. + +## Session-aware routing + +`wait-and-widen` and `pulsar-wait-and-widen` can tell the first request of a +session from its follow-up requests. Set `last_cluster_affinity` to `true` on +the algorithm configuration to enable this. The field defaults to `false`. +Stargate rejects it at startup for every other algorithm, and routing +expressions do not accept it. + +A session is one `x-routing-key`, `x-model`, and `x-cache-affinity-key` tuple. +Stargate keeps no per-session state. The trusted gateway remembers the +`x-stargate-cluster-id` of the session's last 2xx response and sends it on the +next request as `x-stargate-last-cluster-id`. See the +[API gateway contract](api-gateway-contract.md) for the gateway rules. + +This section uses these terms: + +- Affinity order: the ordered cluster list derived from the affinity key. + `wait-and-widen` uses the clusters the hash ring selects, in ring-walk order. + `pulsar-wait-and-widen` uses the full Pulsar ranking. +- Affinity group: the first k clusters of the affinity order, where k is + `cache_affinity_backend_selection_count`. +- X: `cache_affinity_wait_ms`. +- s: `cache_affinity_input_tokens_scale`. + +### Session classification + +When `last_cluster_affinity` is `true`, Stargate classifies each request once, +on its first routing decision. Routing waits and retries keep that +classification. + +| Condition | Session state | +| --- | --- | +| No affinity group: `x-cache-affinity-key` is absent, or `wait-and-widen` has no k | Not applicable. Existing behavior; the header is ignored. | +| Header absent, blank, or not valid UTF-8 | `new` | +| Header names a cluster that is not a current candidate | `stale`, routed as `new` | +| Header names a current candidate | `returning` | + +A current candidate is any routable cluster for the routing key and model, +before retry exclusions. An invalid header value never returns HTTP `400`. +In `pulsar-wait-and-widen`, a hint that names a candidate with no valid Pulsar +weight, and so no place in the ranking, is classified `returning` but cannot +be promoted. The request routes with the original ranking and the configured +X and s. + +The flag belongs to one algorithm configuration. When `x-routing-method` +selects a configuration without it, the header has no effect. + +### Selection rules + +For `new` and `stale` sessions, both algorithms treat X as `0` and s as `1.0` +for the request. A new session has no KV cache to wait for, so it goes to the +best-ranked cluster that is selectable now instead of holding for rank 1. The +affinity order, k, TTFT buckets, queue admission, `band_widen_interval_ms`, and +`fallback_max_queued` apply as configured. When `band_widen_interval_ms` is +unset, it defaults to the effective X of `0`, so every ranked cluster opens as +soon as the affinity group cannot serve. + +For `returning` sessions, both algorithms: + +1. Move the last cluster to position 1 of the affinity order. Every other + cluster keeps its relative order. +2. Take the affinity group from the reordered order. The group keeps size k. + If the last cluster was outside the original group, the original group's + last member drops out. +3. Apply X, s, and every other setting as configured. + +`pulsar-wait-and-widen` cuts its widening bands from the reordered ranking. +`wait-and-widen` global selection is unchanged, because it already covers +every cluster. + +Retry exclusion and KV free-token checks apply to the last cluster as they do +to any other candidate. An excluded or KV-infeasible last cluster behaves like +an ineligible rank-1 cluster. With k = 1, the request waits until X and then +widens. With a larger group, the other group members can still serve the +request during X. The promotion applies to one request. Per-key hash ring selections keep the +original order, so a later request for the same key without the header uses +the original affinity order. + +Example for `pulsar-wait-and-widen` with k = 1, ranking A through E, and last +cluster D: + +```text +original ranking A B C D E +reordered ranking D A B C E +elapsed < X D (affinity group, discounted prefill) +X <= elapsed < X+S D A B +X+S <= elapsed D A B C E +``` + +This keeps a session on the cluster that holds its KV cache. Without the hint, +a session that overflowed to a lower-ranked cluster returns to rank 1 when +rank 1 frees up and loses the cache it built on the overflow cluster. + +Example: + +```json +{ + "default": "power-of-n", + "models": { + "model-a": { + "algorithm": "pulsar-wait-and-widen", + "seed": "model-a-v1", + "require_cache_affinity_key": true, + "max_queued": 2, + "cache_affinity_wait_ms": 300, + "cache_affinity_input_tokens_scale": 0.1, + "band_widen_interval_ms": 100, + "last_cluster_affinity": true + } + } +} +``` + +### Rollout order + +1. Deploy Stargate with `last_cluster_affinity` off. It consumes and ignores + `x-stargate-last-cluster-id`. +2. Deploy the gateway with `STARGATE_LAST_CLUSTER_ENABLED=true`. +3. Turn on `last_cluster_affinity` per model. + +Do not turn on `last_cluster_affinity` before the gateway sends the header. +Otherwise every request is classified `new` and loses its affinity wait and +prefill discount. + ## Algorithm fields `power-of-n` and `wait-and-widen` support this field: @@ -423,6 +565,7 @@ apply only when the selection count is set. | `max_queued` | unsigned integer | `0` | Additional queued requests allowed above `max_engine_concurrency`. A reported concurrency of `0` disables this capacity check. | | `ignore_queue_time` | boolean | `false` | Removes queue delay from TTFT bucket formation. Queue-SLO filtering and the configured comparator are unchanged. | | `ignore_input_processing_time` | boolean | `false` | Removes request prefill time from TTFT bucket formation. The configured comparator is unchanged. | +| `last_cluster_affinity` | boolean | `false` | Classifies requests as new or returning sessions from `x-stargate-last-cluster-id`. See [Session-aware routing](#session-aware-routing). Rejected for other algorithms and in routing expressions. | Stargate does not range-check `next_bucket_unlock_factor`. Values from `0` to `1` unlock a later bucket between no wait and the full TTFT gap. Values outside @@ -502,7 +645,8 @@ x-routing-method: pulsar;seed=stable-a;consider_kv_free_tokens=true ``` Parameter names are the configuration fields above and must apply to the -selected algorithm. Omitted fields keep their configured values. The value is +selected algorithm. `last_cluster_affinity` is not an expression parameter; +set it in the configuration file. Omitted fields keep their configured values. The value is an RFC 8941 Item with Parameters, at most 1024 bytes and 32 parameters, with every parameter written as `key=value` and no commas. Quote a number that needs more than three fractional digits, such as `next_bucket_unlock_factor="0.0625"`. @@ -533,9 +677,11 @@ These proxy headers affect load-balancer behavior: | `x-priority` | optional `u32`, default `0` | Chooses the nearest published queue estimate at or below this priority. | | `x-request-slo-ms` | optional `u64` | Interpolates queue bounds. Does not set or shorten the affinity wait. | | `x-max-wait-ms` | optional `u64` | Routing wait limit from request arrival, capped at 60 seconds. Can expire before global buckets open. | +| `x-stargate-last-cluster-id` | optional | Cluster that served the session's last 2xx response. Used only by algorithms with `last_cluster_affinity`. Blank or invalid UTF-8 means absent. | -Invalid required or numeric values return HTTP `400`. `x-routing-method` is -consumed by Stargate and is not forwarded upstream. See the +Invalid required or numeric values return HTTP `400`. `x-routing-method` and +`x-stargate-last-cluster-id` are consumed by Stargate and are not forwarded +upstream. See the [API gateway contract](api-gateway-contract.md) for the complete proxy header contract. @@ -574,6 +720,19 @@ The existing fallback selection series therefore records global escalation; deployments that previously saw only primary selections can see new fallback counts. Metric names and label sets are unchanged. +With `last_cluster_affinity`, `stargate_routing_session_selections_total` +counts selections by `session_state` (`new`, `returning`, or `stale`) and +`selection`. For `returning`, `primary` means the request went to its last +cluster. For `new` and `stale`, `primary` means the selected cluster came from +the affinity group. The counter starts at zero for every state and selection +pair of a routing key, model, and algorithm once Stargate classifies a request +for them. Requests that are not classified record nothing. The proxy request +span records the state in `routing.session_state`. A `debug` log line records +the state and whether the last cluster was promoted, with the request ID, +model, and routing key. For a promoted `pulsar-wait-and-widen` request, the +rank in `stargate_routing_selections_total` is relative to the reordered +ranking. + ## Validation checklist 1. Parse the JSON before deployment. @@ -587,6 +746,9 @@ counts. Metric names and label sets are unchanged. 6. Send an unconfigured override and confirm HTTP `400`. 7. Exercise an expected fallback and inspect selection, attempt, retry, and exhaustion counters. +8. With `last_cluster_affinity`, send requests with and without + `x-stargate-last-cluster-id` and confirm the session states in + `stargate_routing_session_selections_total`. ## Implementation sources @@ -596,6 +758,7 @@ counts. Metric names and label sets are unchanged. - `crates/stargate/src/load_balancer/wait_and_widen.rs` - `crates/stargate/src/load_balancer/pulsar.rs` - `crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs` +- `crates/stargate/src/load_balancer/session.rs` - `crates/stargate/src/http_proxy/` - `crates/stargate/src/metrics.rs` - `benches/` From 60f81cd3a1ecee9d3f31c1ee8b174546e8ae2214 Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 14:39:20 -0700 Subject: [PATCH 2/6] feat(llm-api-gateway): send each session's last Stargate cluster as a routing hint Remember the x-stargate-cluster-id from each session's last 2xx Stargate response and send it back as x-stargate-last-cluster-id on the session's next request, so Stargate's last_cluster_affinity can tell new sessions from returning ones. - Config, all off by default: STARGATE_LAST_CLUSTER_ENABLED (false), STARGATE_LAST_CLUSTER_TTL (10m), STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT (20ms), STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES (100000). - Only prompt_cache_key, conversation_id, and x-multi-turn-session-id header sessions are eligible. Payload-derived sessions are skipped. - Store: a stargate-last-cluster DMap on the embedded Olric node when OLRIC_ENABLED=true, otherwise an in-process LRU with the same TTL. Keys are lc:v1: plus a SHA-256 of the length-prefixed routing key, model, and cache affinity key. Values over 256 bytes are not stored. - Complete, Stream, and Proxy always strip an inbound x-stargate-last-cluster-id. Lookups run with a hard deadline and never fail the request. Writes run off the request path when 2xx response headers arrive, before the stream ends. - The Olric node now starts whenever OLRIC_ENABLED=true, not only when the rate limiter is on, and is shared by the rate limiter and the last-cluster store. - Add llm_api_gateway_last_cluster_lookups_total, llm_api_gateway_last_cluster_writes_total, and llm_api_gateway_last_cluster_lookup_duration_seconds, pre-initialized to zero, plus last_cluster_lookup and last_cluster_write child spans and warn logs without affinity keys or session IDs. - Expose the four env vars through the llm-api-gateway Helm chart values. - Document the header in the Stargate API gateway contract, the gateway README, the LLM gateway overview, and the gateway metrics reference. No new third-party dependencies. Refs #2289 Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- deploy/helm/llm-api-gateway/README.md | 3 +- .../llm-api-gateway/templates/configmap.yaml | 4 + .../llm-api-gateway/values.yaml | 22 +- .../scripts/test-default-ownership.sh | 14 + .../metrics/llm-api-gateway/metrics.md | 3 + docs/overview/llm-gateway.md | 45 +- .../llm-api-gateway/README.md | 54 +- .../llm-api-gateway/api/BUILD.bazel | 2 + .../llm-api-gateway/api/last_cluster_test.go | 201 +++++++ .../llm-api-gateway/api/messages_handler.go | 3 +- .../llm-api-gateway/api/session_affinity.go | 8 +- .../llm-api-gateway/config/config.go | 51 ++ .../llm-api-gateway/config/config_test.go | 60 ++ .../internal/olrictest/BUILD.bazel | 34 ++ .../internal/olrictest/olrictest.go | 125 ++++ .../llm-api-gateway/lastcluster/BUILD.bazel | 67 +++ .../lastcluster/local_store.go | 112 ++++ .../lastcluster/local_store_test.go | 149 +++++ .../lastcluster/olric_store.go | 59 ++ .../lastcluster/olric_store_test.go | 93 +++ .../llm-api-gateway/lastcluster/store.go | 77 +++ .../llm-api-gateway/lastcluster/tracker.go | 267 +++++++++ .../lastcluster/tracker_test.go | 523 +++++++++++++++++ .../llm-api-gateway/provider/BUILD.bazel | 3 + .../llm-api-gateway/provider/provider.go | 7 + .../llm-api-gateway/provider/stargate.go | 59 ++ .../provider/stargate_last_cluster_test.go | 543 ++++++++++++++++++ .../llm-api-gateway/requestctx/requestctx.go | 9 + .../llm-api-gateway/server/BUILD.bazel | 7 +- .../server/last_cluster_test.go | 177 ++++++ .../llm-api-gateway/server/server.go | 164 +++++- .../llm-api-gateway/telemetry/metrics.go | 81 +++ .../llm-api-gateway/telemetry/metrics_test.go | 95 +++ .../stargate/docs/api-gateway-contract.md | 19 + 34 files changed, 3095 insertions(+), 45 deletions(-) create mode 100644 src/invocation-plane-services/llm-api-gateway/api/last_cluster_test.go create mode 100644 src/invocation-plane-services/llm-api-gateway/internal/olrictest/BUILD.bazel create mode 100644 src/invocation-plane-services/llm-api-gateway/internal/olrictest/olrictest.go create mode 100644 src/invocation-plane-services/llm-api-gateway/lastcluster/BUILD.bazel create mode 100644 src/invocation-plane-services/llm-api-gateway/lastcluster/local_store.go create mode 100644 src/invocation-plane-services/llm-api-gateway/lastcluster/local_store_test.go create mode 100644 src/invocation-plane-services/llm-api-gateway/lastcluster/olric_store.go create mode 100644 src/invocation-plane-services/llm-api-gateway/lastcluster/olric_store_test.go create mode 100644 src/invocation-plane-services/llm-api-gateway/lastcluster/store.go create mode 100644 src/invocation-plane-services/llm-api-gateway/lastcluster/tracker.go create mode 100644 src/invocation-plane-services/llm-api-gateway/lastcluster/tracker_test.go create mode 100644 src/invocation-plane-services/llm-api-gateway/provider/stargate_last_cluster_test.go create mode 100644 src/invocation-plane-services/llm-api-gateway/server/last_cluster_test.go diff --git a/deploy/helm/llm-api-gateway/README.md b/deploy/helm/llm-api-gateway/README.md index 33ac83fcf6..4803c2f4fa 100644 --- a/deploy/helm/llm-api-gateway/README.md +++ b/deploy/helm/llm-api-gateway/README.md @@ -69,7 +69,8 @@ Important settings to review before deployment: - `llmApiGateway.config.nvcfGrpc*` for optional NVCF gRPC auth integration - `llmApiGateway.metrics.enabled` to expose a metrics port on the Service and Deployment (default: `false`) - `llmApiGateway.metrics.serviceMonitor.enabled` to create a Prometheus `ServiceMonitor` (requires `metrics.enabled`) -- `llmApiGateway.olric.*` for embedded rate-limit state and peer discovery +- `llmApiGateway.config.lastCluster.*` for session last-cluster hints sent to the LLM Request Router (default: disabled). Sets `STARGATE_LAST_CLUSTER_ENABLED`, `STARGATE_LAST_CLUSTER_TTL` (`10m`), `STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT` (`20ms`), and `STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES` (`100000`) +- `llmApiGateway.olric.*` for embedded rate-limit and last-cluster state and peer discovery. The Olric node starts whenever `olric.enabled` is `true`, even with rate limiting off - `llmApiGateway.vault.*` for JWT authentication path, role, and audience values used by the Vault Agent injector The default values include development-oriented placeholders. Override them before using the chart in any shared or production environment. diff --git a/deploy/helm/llm-api-gateway/llm-api-gateway/templates/configmap.yaml b/deploy/helm/llm-api-gateway/llm-api-gateway/templates/configmap.yaml index f7a35c102a..71a522320d 100644 --- a/deploy/helm/llm-api-gateway/llm-api-gateway/templates/configmap.yaml +++ b/deploy/helm/llm-api-gateway/llm-api-gateway/templates/configmap.yaml @@ -27,6 +27,10 @@ data: STARGATE_URL: {{ required "llmApiGateway.config.requestRouterUrl is required" $cfg.requestRouterUrl | quote }} STARGATE_CONNECT_TIMEOUT: {{ $cfg.requestRouterConnectTimeout | quote }} STARGATE_REQUEST_TIMEOUT: {{ $cfg.requestRouterRequestTimeout | quote }} + STARGATE_LAST_CLUSTER_ENABLED: {{ $cfg.lastCluster.enabled | toString | quote }} + STARGATE_LAST_CLUSTER_TTL: {{ $cfg.lastCluster.ttl | quote }} + STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT: {{ $cfg.lastCluster.lookupTimeout | quote }} + STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES: {{ $cfg.lastCluster.localMaxEntries | int64 | toString | quote }} NVCF_GATEWAY_INFERENCE_WRITE_TIMEOUT: {{ $cfg.inferenceWriteTimeout | quote }} NVCF_GRPC_ADDR: {{ $cfg.nvcfGrpcAddr | quote }} NVCF_GRPC_INSECURE: {{ $cfg.nvcfGrpcInsecure | toString | quote }} diff --git a/deploy/helm/llm-api-gateway/llm-api-gateway/values.yaml b/deploy/helm/llm-api-gateway/llm-api-gateway/values.yaml index 5594036e34..05d18443f2 100644 --- a/deploy/helm/llm-api-gateway/llm-api-gateway/values.yaml +++ b/deploy/helm/llm-api-gateway/llm-api-gateway/values.yaml @@ -79,11 +79,23 @@ llmApiGateway: defaultModel: "" defaultRpm: 120 defaultTpm: 240000 - - # Rate-limit state is kept in an embedded Olric cluster colocated with the - # gateway pods (no separate Redis deployment). The gateway's server.go picks - # ratelimit.AllowAll when RATE_LIMIT_ENABLED=true AND OLRIC_ENABLED=false, so - # enabling rate limiting without Olric silently turns into a no-op. + # Session last-cluster hint sent to the LLM Request Router as + # x-stargate-last-cluster-id. Entries live in the embedded Olric cluster + # when olric.enabled is true, otherwise in a per-pod LRU capped at + # localMaxEntries. Enable it after the router is upgraded and before any + # model turns on last_cluster_affinity. + lastCluster: + enabled: false + ttl: 10m + lookupTimeout: 20ms + localMaxEntries: 100000 + + # Rate-limit state and session last-cluster hints are kept in an embedded + # Olric cluster colocated with the gateway pods (no separate Redis + # deployment). The node starts whenever olric.enabled is true. The gateway's + # server.go picks ratelimit.AllowAll when RATE_LIMIT_ENABLED=true AND + # OLRIC_ENABLED=false, so enabling rate limiting without Olric silently turns + # into a no-op. olric: enabled: true # "local" | "lan" | "wan" - memberlist tuning. Use "local" for single-pod k3d diff --git a/deploy/helm/llm-api-gateway/scripts/test-default-ownership.sh b/deploy/helm/llm-api-gateway/scripts/test-default-ownership.sh index 7a1abc4431..2fe99dce43 100755 --- a/deploy/helm/llm-api-gateway/scripts/test-default-ownership.sh +++ b/deploy/helm/llm-api-gateway/scripts/test-default-ownership.sh @@ -32,6 +32,14 @@ test "$(yq ea -r 'select(.kind == "Deployment") | [.spec.template.spec.container fail "generic chart must leave the metrics endpoint disabled by default" test "$(yq ea -r '[select(.kind == "ServiceMonitor")] | length' "$default_manifest")" = 0 || fail "generic chart must not create a ServiceMonitor by default" +test "$(read_config "$default_manifest" STARGATE_LAST_CLUSTER_ENABLED)" = false || + fail "generic chart must leave last-cluster hints disabled by default" +test "$(read_config "$default_manifest" STARGATE_LAST_CLUSTER_TTL)" = 10m || + fail "last-cluster TTL default must be 10m" +test "$(read_config "$default_manifest" STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT)" = 20ms || + fail "last-cluster lookup timeout default must be 20ms" +test "$(read_config "$default_manifest" STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES)" = 100000 || + fail "last-cluster local store cap default must be 100000" helm template llm-api-gateway "$chart_dir" \ --namespace nvcf \ @@ -39,10 +47,16 @@ helm template llm-api-gateway "$chart_dir" \ --set llmApiGateway.config.nvcfGrpcInsecure=true \ --set llmApiGateway.metrics.enabled=true \ --set llmApiGateway.metrics.serviceMonitor.enabled=true \ + --set llmApiGateway.config.lastCluster.enabled=true \ + --set llmApiGateway.config.lastCluster.localMaxEntries=1000000 \ >"$override_manifest" test "$(read_config "$override_manifest" NVCF_GRPC_INSECURE)" = true || fail "explicit plaintext transport override did not reach the ConfigMap" +test "$(read_config "$override_manifest" STARGATE_LAST_CLUSTER_ENABLED)" = true || + fail "last-cluster opt-in did not reach the ConfigMap" +test "$(read_config "$override_manifest" STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES)" = 1000000 || + fail "large last-cluster cap must render as a plain integer" test "$(yq ea -r 'select(.kind == "Deployment") | [.spec.template.spec.containers[0].ports[] | select(.name == "metrics")] | length' "$override_manifest")" = 1 || fail "explicit metrics override did not expose the metrics container port" test "$(yq ea -r '[select(.kind == "ServiceMonitor" and .metadata.name == "llm-api-gateway-metrics")] | length' "$override_manifest")" = 1 || diff --git a/docs/observability/metrics/llm-api-gateway/metrics.md b/docs/observability/metrics/llm-api-gateway/metrics.md index 49dcb8f3e8..36920f9273 100644 --- a/docs/observability/metrics/llm-api-gateway/metrics.md +++ b/docs/observability/metrics/llm-api-gateway/metrics.md @@ -40,3 +40,6 @@ other unbounded request fields as metric labels. | `llm_api_gateway_rate_limit_synchronizer_queue_wait_seconds` | Histogram | `llm-api-gateway:9464/metrics` | None | Time spent queueing a rate limit event in seconds. | | `llm_api_gateway_rate_limit_synchronizer_queue_length` | Gauge | `llm-api-gateway:9464/metrics` | None | Current rate limit synchronizer queue length. | | `llm_api_gateway_rate_limit_synchronizer_events_dropped_total` | Counter | `llm-api-gateway:9464/metrics` | `reason` | Number of rate limit events dropped before publishing. `reason` is a bounded enum such as `old_message`. | +| `llm_api_gateway_last_cluster_lookups_total` | Counter | `llm-api-gateway:9464/metrics` | `result` | Session last-cluster store lookups. `result` is `hit`, `miss`, `error`, `timeout`, or `skipped`. All values start at zero. | +| `llm_api_gateway_last_cluster_writes_total` | Counter | `llm-api-gateway:9464/metrics` | `result` | Session last-cluster store writes. `result` is `ok` or `error`. All values start at zero. | +| `llm_api_gateway_last_cluster_lookup_duration_seconds` | Histogram | `llm-api-gateway:9464/metrics` | `result` | Session last-cluster store lookup duration in seconds, with buckets from 1 ms to 50 ms. Skipped lookups are not timed. | diff --git a/docs/overview/llm-gateway.md b/docs/overview/llm-gateway.md index 63b8d3c225..88ab44c575 100644 --- a/docs/overview/llm-gateway.md +++ b/docs/overview/llm-gateway.md @@ -419,6 +419,46 @@ sequenceDiagram Sticky routing only affects backend selection when the LLM request router is configured with a cache-affinity-aware routing method for the target model. +### Last-Cluster Hints + +The gateway can remember which LLM request router cluster served each +session's last successful request and send it back on the next request in the +trusted `x-stargate-last-cluster-id` header. The router uses the hint when the +target model's `wait-and-widen` or `pulsar-wait-and-widen` configuration sets +`last_cluster_affinity: true`. A returning session then waits for the cluster +that holds its KV cache, and a new session goes straight to the best cluster +that is free now. + +A session is the routing key, model, and cache affinity key. Only sessions +from `prompt_cache_key`, `conversation.id`, or `x-multi-turn-session-id` are +eligible. Sessions that fall back to a messages or input hash change key every +turn, so the gateway never looks them up or stores them. + +The gateway records the `x-stargate-cluster-id` of each 2xx router response +when the response headers arrive. Errors never change the stored value. A +client-supplied `x-stargate-last-cluster-id` is always stripped. + +| Env var | Helm value | Default | Effect | +| --- | --- | --- | --- | +| `STARGATE_LAST_CLUSTER_ENABLED` | `llmApiGateway.config.lastCluster.enabled` | `false` | Turns lookup, the header, and store writes on. | +| `STARGATE_LAST_CLUSTER_TTL` | `llmApiGateway.config.lastCluster.ttl` | `10m` | Entry lifetime, reset on every write. | +| `STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT` | `llmApiGateway.config.lastCluster.lookupTimeout` | `20ms` | Lookup deadline on the request path. The request proceeds without the hint after it. | +| `STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES` | `llmApiGateway.config.lastCluster.localMaxEntries` | `100000` | Per-replica LRU cap when Olric is off. | + +With `OLRIC_ENABLED=true`, entries live in the `stargate-last-cluster` DMap of +the gateway's embedded Olric cluster and are shared by all replicas. The Olric +node starts whenever Olric is enabled, even with rate limiting off. + +Roll out in this order: + +1. Upgrade the LLM request router with `last_cluster_affinity` off. It accepts + the header and ignores it. +2. Set `STARGATE_LAST_CLUSTER_ENABLED=true` on the gateway. +3. Turn on `last_cluster_affinity` per model in the router configuration. + +Do not turn on `last_cluster_affinity` before step 2. Without the hint, every +request is treated as a new session and skips the affinity wait. + ## Metrics LLM API Gateway request metrics include a `function_id` label. The value is the @@ -437,9 +477,12 @@ key, such as health checks, use `function_id="none"`. | `llm_api_gateway_provider_time_seconds` | Histogram | `endpoint`, `phase`, `stream`, `function_id` | Provider-reported timing phases. | | `llm_api_gateway_stream_first_token_seconds` | Histogram | `endpoint`, `function_id` | Time from stream request start to the first token. | | `llm_api_gateway_stream_duration_seconds` | Histogram | `endpoint`, `status`, `function_id` | Total stream duration. | +| `llm_api_gateway_last_cluster_lookups_total` | Counter | `result` | Last-cluster store lookups: `hit`, `miss`, `error`, `timeout`, or `skipped`. | +| `llm_api_gateway_last_cluster_writes_total` | Counter | `result` | Last-cluster store writes: `ok` or `error`. | +| `llm_api_gateway_last_cluster_lookup_duration_seconds` | Histogram | `result` | Last-cluster store lookup latency, 1 ms to 50 ms buckets. | Infrastructure metrics for authentication, rate limit synchronization, pub/sub, -and the distributed cache do not include `function_id` because they are not +the last-cluster store, and the distributed cache do not include `function_id` because they are not associated with a single routed function. Use the label to calculate request rate by function: diff --git a/src/invocation-plane-services/llm-api-gateway/README.md b/src/invocation-plane-services/llm-api-gateway/README.md index 961509588d..8ab3c28507 100644 --- a/src/invocation-plane-services/llm-api-gateway/README.md +++ b/src/invocation-plane-services/llm-api-gateway/README.md @@ -72,6 +72,37 @@ conversation. The gateway preserves the raw body field for the model backend. It forwards only a SHA-256-derived value in the internal `x-cache-affinity-key` header to Stargate. +### Last-cluster hints + +With `STARGATE_LAST_CLUSTER_ENABLED=true`, the gateway remembers which +Stargate cluster served each session's last successful request and sends it +back on the session's next request as `x-stargate-last-cluster-id`. Stargate +uses the hint only for models whose routing algorithm sets +`last_cluster_affinity`. + +- A session is the tuple of routing key, model, and cache affinity key. +- Only sessions from `prompt_cache_key`, `conversation_id`, or the + `x-multi-turn-session-id` header are eligible. Payload-derived sessions hash + the full message list, which changes every turn, so they are never looked up + or stored. Sessions from `x-claude-code-session-id` are also skipped. +- The gateway writes the `x-stargate-cluster-id` of a 2xx Stargate response + when the response headers arrive, before a stream ends. Non-2xx responses and + transport errors never change the stored value. Writes run in the background + and never delay or fail the client response. +- The lookup is bounded by `STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT`. On a miss, + error, or timeout the request proceeds without the hint. +- Store keys are `lc:v1:` plus the hex SHA-256 of the length-prefixed routing + key, model, and cache affinity key. Raw routing keys never appear in store + keys. Affinity keys and session IDs never appear in store keys or logs. + Cluster IDs over 256 bytes are not stored. +- With `OLRIC_ENABLED=true`, entries live in the `stargate-last-cluster` DMap + of the embedded Olric cluster, so every replica sees them. With Olric off, + each replica keeps its own LRU capped at + `STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES`. Every write resets the entry TTL. +- The gateway always strips a client-supplied `x-stargate-last-cluster-id`, + whether or not the feature is enabled, and keeps stripping `x-stargate-*` + response headers from client responses. + When `NVCF_GRPC_ADDR` is configured, the gateway authenticates each request through the NVCF LLM gRPC auth service, derives the per-caller rate-limit key from `authContext["ncaId"]`, optionally scopes it further by project, and keeps @@ -146,6 +177,14 @@ Useful overrides: 413 (default `0`, no limit) - `STARGATE_CONNECT_TIMEOUT` to control Stargate dial timeout - `STARGATE_REQUEST_TIMEOUT` to cap end-to-end Stargate request time +- `STARGATE_LAST_CLUSTER_ENABLED=true` to send session last-cluster hints to + Stargate (default `false`) +- `STARGATE_LAST_CLUSTER_TTL` for the last-cluster entry lifetime, reset on + every write (default `10m`) +- `STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT` for the lookup deadline on the + request path (default `20ms`) +- `STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES` for the per-replica store size + when Olric is off (default `100000`) - `NVCF_GATEWAY_INFERENCE_WRITE_TIMEOUT` to cap how long one response write may stall on a client that stopped reading (default `60s`, `0s` disables). It applies only while a write is in progress, so long streams, long @@ -160,7 +199,9 @@ Useful overrides: - `NVCF_GRPC_TIMEOUT` to cap each gRPC auth or policy call - `RATE_LIMIT_ENABLED=false` to disable rate limiting locally - `RATE_LIMIT_FAIL_OPEN=false` to make Olric or limiter failures fatal -- `OLRIC_ENABLED=false` to skip starting the embedded Olric node +- `OLRIC_ENABLED=false` to skip starting the embedded Olric node. When it is + `true`, the node starts even if rate limiting is off, and it also holds + last-cluster hints - `OLRIC_BIND_PORT`, `OLRIC_MEMBERLIST_BIND_PORT`, and `OLRIC_PEERS` for multi-instance Olric clustering - `OTEL_SERVICE_NAME` to override the emitted service name @@ -178,6 +219,17 @@ time, first-token time, and stream duration metrics. Infrastructure metrics for authentication, pub/sub, rate-limit synchronization, and Olric remain function-independent. +Last-cluster store metrics use only a `result` label: + +- `llm_api_gateway_last_cluster_lookups_total`: `hit`, `miss`, `error`, + `timeout`, or `skipped` (enabled, but the session is not eligible) +- `llm_api_gateway_last_cluster_writes_total`: `ok` or `error` +- `llm_api_gateway_last_cluster_lookup_duration_seconds`: lookup latency, with + buckets from 1 ms to 50 ms + +Lookups and writes also create the `llm-api-gateway.last_cluster_lookup` and +`llm-api-gateway.last_cluster_write` spans. + Example request-rate query: ```promql diff --git a/src/invocation-plane-services/llm-api-gateway/api/BUILD.bazel b/src/invocation-plane-services/llm-api-gateway/api/BUILD.bazel index b0514ee497..922a4e3b2b 100644 --- a/src/invocation-plane-services/llm-api-gateway/api/BUILD.bazel +++ b/src/invocation-plane-services/llm-api-gateway/api/BUILD.bazel @@ -89,6 +89,7 @@ go_test( "body_limit_test.go", "embeddings_handler_test.go", "info_test.go", + "last_cluster_test.go", "middleware_telemetry_test.go", "messages_handler_test.go", "middleware_test.go", @@ -113,6 +114,7 @@ go_test( "//src/invocation-plane-services/llm-api-gateway/api/adapters/openairesponses", "//src/invocation-plane-services/llm-api-gateway/config", "//src/invocation-plane-services/llm-api-gateway/internal/ptr", + "//src/invocation-plane-services/llm-api-gateway/lastcluster", "//src/invocation-plane-services/llm-api-gateway/models", "//src/invocation-plane-services/llm-api-gateway/nvcf", "//src/invocation-plane-services/llm-api-gateway/provider", diff --git a/src/invocation-plane-services/llm-api-gateway/api/last_cluster_test.go b/src/invocation-plane-services/llm-api-gateway/api/last_cluster_test.go new file mode 100644 index 0000000000..f8b959c0af --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/api/last_cluster_test.go @@ -0,0 +1,201 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package api + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + echo "github.com/labstack/echo/v4" + + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/config" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/lastcluster" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/provider" +) + +// lastClusterAPI serves the gateway routes against a fake Stargate that sets +// x-stargate-cluster-id and records the last-cluster hint per request. +type lastClusterAPI struct { + e *echo.Echo + tracker *lastcluster.Tracker + store *lastcluster.LocalStore + mu sync.Mutex + hints []string +} + +func newLastClusterAPI(t *testing.T) *lastClusterAPI { + t.Helper() + gw := &lastClusterAPI{} + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + gw.mu.Lock() + gw.hints = append(gw.hints, r.Header.Get("X-Stargate-Last-Cluster-Id")) + gw.mu.Unlock() + w.Header().Set("X-Stargate-Cluster-Id", "cluster-b") + w.Header().Set(echo.HeaderContentType, "text/event-stream") + w.WriteHeader(http.StatusOK) + switch r.URL.Path { + case "/v1/responses": + _, _ = io.WriteString(w, "event: response.completed\n"+ + `data: {"type":"response.completed","sequence_number":0,"response":{"id":"resp_1","object":"response","status":"completed","created_at":1,"model":"company-name/model-name","output":[]}}`+"\n\n") + default: + _, _ = io.WriteString(w, `data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"company-name/model-name","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]}`+"\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + } + })) + t.Cleanup(upstream.Close) + + cfg := config.Default() + p, err := provider.NewStargateProvider(config.StargateConfig{URL: upstream.URL}) + if err != nil { + t.Fatal(err) + } + gw.store = lastcluster.NewLocalStore(100) + gw.tracker = lastcluster.NewTracker(gw.store, lastcluster.Options{TTL: time.Minute, LookupTimeout: 20 * time.Millisecond}) + p.SetLastClusterTracker(gw.tracker) + gw.e = echo.New() + gw.e.Use(NewContextMiddleware(cfg)) + RegisterRoutes(gw.e, NewHandlers(cfg, p, nil, nil)) + return gw +} + +func (a *lastClusterAPI) post(t *testing.T, path, payload string, header http.Header) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(payload)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + for name, values := range header { + for _, value := range values { + req.Header.Add(name, value) + } + } + rec := httptest.NewRecorder() + a.e.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("%s status = %d: %s", path, rec.Code, rec.Body.String()) + } + a.tracker.Drain(context.Background()) + return rec +} + +func (a *lastClusterAPI) lastHint() string { + a.mu.Lock() + defer a.mu.Unlock() + return a.hints[len(a.hints)-1] +} + +func TestLastClusterSessionFlows(t *testing.T) { + sessionHeader := http.Header{HeaderMultiTurnSessionID: []string{"client-session-1"}} + for _, tc := range []struct { + name string + path string + payload string + header http.Header + }{ + { + name: "chat/prompt_cache_key", + path: "/v1/chat/completions", + payload: `{"model":"fn-alpha/company-name/model-name","prompt_cache_key":"pck-1","messages":[{"role":"user","content":"hello"}]}`, + }, + { + name: "chat/header", + path: "/v1/chat/completions", + payload: `{"model":"fn-alpha/company-name/model-name","stream":true,"messages":[{"role":"user","content":"hello"}]}`, + header: sessionHeader, + }, + { + name: "responses/prompt_cache_key", + path: "/v1/responses", + payload: `{"model":"fn-alpha/company-name/model-name","stream":true,"prompt_cache_key":"pck-1","input":"hello"}`, + }, + { + name: "responses/conversation_id", + path: "/v1/responses", + payload: `{"model":"fn-alpha/company-name/model-name","stream":true,"conversation":{"id":"conv-1"},"input":"hello"}`, + }, + { + name: "responses/header", + path: "/v1/responses", + payload: `{"model":"fn-alpha/company-name/model-name","stream":true,"input":"hello"}`, + header: sessionHeader, + }, + } { + t.Run(tc.name, func(t *testing.T) { + gw := newLastClusterAPI(t) + + gw.post(t, tc.path, tc.payload, tc.header) + if hint := gw.lastHint(); hint != "" { + t.Fatalf("first request hint = %q, want none", hint) + } + rec := gw.post(t, tc.path, tc.payload, tc.header) + if hint := gw.lastHint(); hint != "cluster-b" { + t.Fatalf("second request hint = %q, want cluster-b", hint) + } + if got := rec.Header().Get("X-Stargate-Cluster-Id"); got != "" { + t.Fatalf("x-stargate-cluster-id leaked to the client: %q", got) + } + }) + } +} + +func TestLastClusterPayloadSessionNeverStored(t *testing.T) { + for _, tc := range []struct{ path, payload string }{ + {"/v1/chat/completions", `{"model":"fn-alpha/company-name/model-name","messages":[{"role":"user","content":"hello"}]}`}, + {"/v1/responses", `{"model":"fn-alpha/company-name/model-name","stream":true,"input":"hello"}`}, + } { + t.Run(tc.path, func(t *testing.T) { + gw := newLastClusterAPI(t) + for i := 0; i < 2; i++ { + gw.post(t, tc.path, tc.payload, nil) + if hint := gw.lastHint(); hint != "" { + t.Fatalf("request %d hint = %q, want none", i, hint) + } + } + if n := gw.store.Len(); n != 0 { + t.Fatalf("store entries = %d, want 0 for payload-derived sessions", n) + } + }) + } +} + +func TestClientLastClusterHeaderNeverForwarded(t *testing.T) { + forged := http.Header{"X-Stargate-Last-Cluster-Id": []string{"client-forged"}} + for _, tc := range []struct{ path, payload string }{ + {"/v1/chat/completions", `{"model":"fn-alpha/company-name/model-name","messages":[{"role":"user","content":"hello"}]}`}, + {"/v1/chat/completions", `{"model":"fn-alpha/company-name/model-name","stream":true,"prompt_cache_key":"pck-2","messages":[{"role":"user","content":"hello"}]}`}, + {"/v1/responses", `{"model":"fn-alpha/company-name/model-name","stream":true,"input":"hello"}`}, + {"/v1/responses", `{"model":"fn-alpha/company-name/model-name","stream":true,"prompt_cache_key":"pck-2","input":"hello"}`}, + } { + t.Run(tc.path, func(t *testing.T) { + gw := newLastClusterAPI(t) + gw.post(t, tc.path, tc.payload, forged) + if hint := gw.lastHint(); hint != "" { + t.Fatalf("first request hint = %q, want the forged value dropped", hint) + } + gw.post(t, tc.path, tc.payload, forged) + if hint := gw.lastHint(); hint == "client-forged" { + t.Fatal("client-supplied hint reached Stargate") + } + }) + } +} diff --git a/src/invocation-plane-services/llm-api-gateway/api/messages_handler.go b/src/invocation-plane-services/llm-api-gateway/api/messages_handler.go index d4a50f046b..d97b738711 100644 --- a/src/invocation-plane-services/llm-api-gateway/api/messages_handler.go +++ b/src/invocation-plane-services/llm-api-gateway/api/messages_handler.go @@ -23,6 +23,7 @@ import ( "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/models" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/provider" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/ratelimit" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/requestctx" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/telemetry" ) @@ -88,7 +89,7 @@ func (h *MessagesHandlers) ServeMessages(ec echo.Context) error { return err } } else if sessionID := c.Request().Header.Get("x-claude-code-session-id"); sessionID != "" { - if err := setSessionAffinity(reqCtx, "claude_code_header", sessionID); err != nil { + if err := setSessionAffinity(reqCtx, requestctx.SessionSourceClaudeCodeHeader, sessionID); err != nil { return err } } else { diff --git a/src/invocation-plane-services/llm-api-gateway/api/session_affinity.go b/src/invocation-plane-services/llm-api-gateway/api/session_affinity.go index 7e904651ad..461bac48ba 100644 --- a/src/invocation-plane-services/llm-api-gateway/api/session_affinity.go +++ b/src/invocation-plane-services/llm-api-gateway/api/session_affinity.go @@ -36,10 +36,10 @@ const ( HeaderMultiTurnSessionID = "x-multi-turn-session-id" sessionIDMaxLen = 256 - sessionAffinitySourceHeader = "header" - sessionAffinitySourcePayload = "payload" - sessionAffinitySourcePrompt = "prompt_cache_key" - sessionAffinitySourceConv = "conversation_id" + sessionAffinitySourceHeader = requestctx.SessionSourceHeader + sessionAffinitySourcePayload = requestctx.SessionSourcePayload + sessionAffinitySourcePrompt = requestctx.SessionSourcePromptCacheKey + sessionAffinitySourceConv = requestctx.SessionSourceConversationID sessionAffinitySourceSession = "session" ) diff --git a/src/invocation-plane-services/llm-api-gateway/config/config.go b/src/invocation-plane-services/llm-api-gateway/config/config.go index cfe7f542f4..3767743231 100644 --- a/src/invocation-plane-services/llm-api-gateway/config/config.go +++ b/src/invocation-plane-services/llm-api-gateway/config/config.go @@ -77,6 +77,20 @@ type StargateConfig struct { URL string ConnectTimeout time.Duration RequestTimeout time.Duration + LastCluster LastClusterConfig +} + +// LastClusterConfig controls the per-session last-cluster hint sent to +// Stargate as x-stargate-last-cluster-id. The store is the shared Olric DMap +// when Olric is enabled, otherwise an in-process LRU capped at LocalMaxEntries. +type LastClusterConfig struct { + Enabled bool + // TTL is the entry lifetime. Every write resets it. + TTL time.Duration + // LookupTimeout bounds the store lookup on the request path. A slower + // lookup is abandoned and the request proceeds without the hint. + LookupTimeout time.Duration + LocalMaxEntries int } type NVCFConfig struct { @@ -210,6 +224,12 @@ func Default() *Config { URL: "http://127.0.0.1:8000", ConnectTimeout: 2 * time.Second, RequestTimeout: 0, + LastCluster: LastClusterConfig{ + Enabled: false, + TTL: 10 * time.Minute, + LookupTimeout: 20 * time.Millisecond, + LocalMaxEntries: 100000, + }, }, NVCF: NVCFConfig{ GRPCAddr: "", @@ -251,6 +271,7 @@ func LoadFromEnv() (*Config, error) { applyServerTelemetryEnv(cfg, &errs) applyStargateNVCFEnv(cfg, &errs) + applyStargateLastClusterEnv(cfg, &errs) applyOlricNetworkEnv(cfg, &errs) applyOlricRuntimeEnv(cfg, &errs) applyRateLimitEnv(cfg, &errs) @@ -362,6 +383,36 @@ func applyStargateNVCFEnv(cfg *Config, errs *envErrs) { } +func applyStargateLastClusterEnv(cfg *Config, errs *envErrs) { + if v, ok := errs.boolean("STARGATE_LAST_CLUSTER_ENABLED"); ok { + cfg.Stargate.LastCluster.Enabled = v + } + + if ttl, ok := errs.duration("STARGATE_LAST_CLUSTER_TTL"); ok { + if ttl <= 0 { + errs.add("STARGATE_LAST_CLUSTER_TTL", ttl.String(), errors.New("must be > 0")) + } else { + cfg.Stargate.LastCluster.TTL = ttl + } + } + + if timeout, ok := errs.duration("STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT"); ok { + if timeout <= 0 { + errs.add("STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT", timeout.String(), errors.New("must be > 0")) + } else { + cfg.Stargate.LastCluster.LookupTimeout = timeout + } + } + + if maxEntries, ok := errs.integer("STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES"); ok { + if maxEntries <= 0 { + errs.add("STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES", strconv.Itoa(maxEntries), errors.New("must be > 0")) + } else { + cfg.Stargate.LastCluster.LocalMaxEntries = maxEntries + } + } +} + func applyOlricNetworkEnv(cfg *Config, errs *envErrs) { if v, ok := errs.boolean("OLRIC_ENABLED"); ok { cfg.Olric.Enabled = v diff --git a/src/invocation-plane-services/llm-api-gateway/config/config_test.go b/src/invocation-plane-services/llm-api-gateway/config/config_test.go index 6ae014ba0f..b40ebe3302 100644 --- a/src/invocation-plane-services/llm-api-gateway/config/config_test.go +++ b/src/invocation-plane-services/llm-api-gateway/config/config_test.go @@ -22,6 +22,7 @@ import ( "path/filepath" "strings" "testing" + "time" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/internal/servicetier" ) @@ -47,6 +48,65 @@ func TestDefaultSetsExpectedDefaults(t *testing.T) { if cfg.Server.InferenceWriteTimeout <= 0 { t.Fatalf("inference write timeout = %s, want a positive stall bound", cfg.Server.InferenceWriteTimeout) } + want := LastClusterConfig{ + Enabled: false, + TTL: 10 * time.Minute, + LookupTimeout: 20 * time.Millisecond, + LocalMaxEntries: 100000, + } + if cfg.Stargate.LastCluster != want { + t.Fatalf("last cluster = %+v, want %+v", cfg.Stargate.LastCluster, want) + } +} + +func TestLoadFromEnvReadsLastCluster(t *testing.T) { + t.Setenv("STARGATE_LAST_CLUSTER_ENABLED", "true") + t.Setenv("STARGATE_LAST_CLUSTER_TTL", "90s") + t.Setenv("STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT", "5ms") + t.Setenv("STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES", "42") + + cfg, err := LoadFromEnv() + if err != nil { + t.Fatalf("LoadFromEnv() error = %v", err) + } + + want := LastClusterConfig{ + Enabled: true, + TTL: 90 * time.Second, + LookupTimeout: 5 * time.Millisecond, + LocalMaxEntries: 42, + } + if cfg.Stargate.LastCluster != want { + t.Fatalf("last cluster = %+v, want %+v", cfg.Stargate.LastCluster, want) + } +} + +func TestLoadFromEnvRejectsInvalidLastCluster(t *testing.T) { + cases := []struct { + key string + value string + }{ + {key: "STARGATE_LAST_CLUSTER_ENABLED", value: "maybe"}, + {key: "STARGATE_LAST_CLUSTER_TTL", value: "10"}, + {key: "STARGATE_LAST_CLUSTER_TTL", value: "0s"}, + {key: "STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT", value: "-1ms"}, + {key: "STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES", value: "0"}, + {key: "STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES", value: "-1"}, + {key: "STARGATE_LAST_CLUSTER_LOCAL_MAX_ENTRIES", value: "many"}, + } + for _, tc := range cases { + t.Run(tc.key+"="+tc.value, func(t *testing.T) { + t.Setenv(tc.key, tc.value) + + _, err := LoadFromEnv() + if err == nil { + t.Fatal("LoadFromEnv() error = nil, want error") + } + if !strings.Contains(err.Error(), tc.key) { + t.Fatalf("error %q does not mention %s", err.Error(), tc.key) + } + }) + } } func TestLoadFromEnvReadsInferenceWriteTimeout(t *testing.T) { diff --git a/src/invocation-plane-services/llm-api-gateway/internal/olrictest/BUILD.bazel b/src/invocation-plane-services/llm-api-gateway/internal/olrictest/BUILD.bazel new file mode 100644 index 0000000000..c845f09780 --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/internal/olrictest/BUILD.bazel @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "olrictest", + testonly = True, + srcs = ["olrictest.go"], + importpath = "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/internal/olrictest", + visibility = ["//src/invocation-plane-services/llm-api-gateway:__subpackages__"], + deps = [ + "//src/invocation-plane-services/llm-api-gateway/config", + "//src/invocation-plane-services/llm-api-gateway/util", + ], +) + +alias( + name = "go_default_library", + actual = ":olrictest", + visibility = ["//src/invocation-plane-services/llm-api-gateway:__subpackages__"], +) diff --git a/src/invocation-plane-services/llm-api-gateway/internal/olrictest/olrictest.go b/src/invocation-plane-services/llm-api-gateway/internal/olrictest/olrictest.go new file mode 100644 index 0000000000..b7cabb5530 --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/internal/olrictest/olrictest.go @@ -0,0 +1,125 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package olrictest starts embedded Olric clusters for tests. +package olrictest + +import ( + "context" + "fmt" + "io" + "net" + "strconv" + "testing" + "time" + + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/config" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/util" +) + +// FreeTCPPort returns a port that was free on 127.0.0.1 when checked. +func FreeTCPPort() (int, error) { + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return 0, err + } + defer l.Close() + addr, ok := l.Addr().(*net.TCPAddr) + if !ok { + return 0, fmt.Errorf("unexpected listener addr type %T", l.Addr()) + } + return addr.Port, nil +} + +// NodeConfig returns a loopback Olric config on fresh ports that joins peers. +func NodeConfig(t testing.TB, peers []string) config.OlricConfig { + t.Helper() + clientPort, err := FreeTCPPort() + if err != nil { + t.Fatalf("free client port: %v", err) + } + memberlistPort, err := FreeTCPPort() + if err != nil { + t.Fatalf("free memberlist port: %v", err) + } + return config.OlricConfig{ + Enabled: true, + Environment: "local", + BindAddr: "127.0.0.1", + BindPort: clientPort, + MemberlistBindAddr: "127.0.0.1", + MemberlistBindPort: memberlistPort, + Peers: peers, + ReplicaCount: 1, + PartitionCount: 7, + DMapName: "olrictest", + StartupTimeout: 15 * time.Second, + ShutdownTimeout: 5 * time.Second, + LogLevel: "ERROR", + LogOutput: io.Discard, + } +} + +// StartCluster boots size embedded nodes joined over memberlist and waits +// until every node sees every member. Nodes shut down on test cleanup. +func StartCluster(t testing.TB, size int) []*util.OlricNode { + t.Helper() + ctx := context.Background() + var ( + nodes []*util.OlricNode + peers []string + ) + for i := 0; i < size; i++ { + cfg := NodeConfig(t, append([]string(nil), peers...)) + node, err := util.NewOlricNode(ctx, cfg) + if err != nil { + t.Fatalf("start olric node %d: %v", i, err) + } + t.Cleanup(func() { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = node.Shutdown(shutdownCtx) + }) + nodes = append(nodes, node) + peers = append(peers, net.JoinHostPort(cfg.MemberlistBindAddr, strconv.Itoa(cfg.MemberlistBindPort))) + } + waitForClusterSize(t, nodes, size, 10*time.Second) + return nodes +} + +func waitForClusterSize(t testing.TB, nodes []*util.OlricNode, want int, timeout time.Duration) { + t.Helper() + ctx := context.Background() + deadline := time.Now().Add(timeout) + for { + ready := true + for _, n := range nodes { + members, err := n.Client.Members(ctx) + if err != nil || len(members) < want { + ready = false + break + } + } + if ready { + return + } + if time.Now().After(deadline) { + t.Fatalf("olric cluster did not reach %d members within %s", want, timeout) + } + time.Sleep(50 * time.Millisecond) + } +} diff --git a/src/invocation-plane-services/llm-api-gateway/lastcluster/BUILD.bazel b/src/invocation-plane-services/llm-api-gateway/lastcluster/BUILD.bazel new file mode 100644 index 0000000000..ec5f563cba --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/lastcluster/BUILD.bazel @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "lastcluster", + srcs = [ + "local_store.go", + "olric_store.go", + "store.go", + "tracker.go", + ], + importpath = "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/lastcluster", + visibility = ["//visibility:public"], + deps = [ + "//src/invocation-plane-services/llm-api-gateway/requestctx", + "//src/invocation-plane-services/llm-api-gateway/telemetry", + "@com_github_olric_data_olric//:olric", + "@com_github_rs_zerolog//:zerolog", + "@io_opentelemetry_go_otel//attribute", + "@io_opentelemetry_go_otel//codes", + "@io_opentelemetry_go_otel_trace//:trace", + ], +) + +alias( + name = "go_default_library", + actual = ":lastcluster", + visibility = ["//visibility:public"], +) + +go_test( + name = "lastcluster_test", + srcs = [ + "local_store_test.go", + "olric_store_test.go", + "tracker_test.go", + ], + embed = [":lastcluster"], + deps = [ + "//src/invocation-plane-services/llm-api-gateway/internal/olrictest", + "//src/invocation-plane-services/llm-api-gateway/requestctx", + "//src/invocation-plane-services/llm-api-gateway/telemetry", + "@com_github_rs_zerolog//:zerolog", + "@io_opentelemetry_go_otel//:otel", + "@io_opentelemetry_go_otel//attribute", + "@io_opentelemetry_go_otel//codes", + "@io_opentelemetry_go_otel_sdk//trace", + "@io_opentelemetry_go_otel_sdk//trace/tracetest", + "@io_opentelemetry_go_otel_sdk_metric//:metric", + "@io_opentelemetry_go_otel_sdk_metric//metricdata", + "@io_opentelemetry_go_otel_trace//:trace", + ], +) diff --git a/src/invocation-plane-services/llm-api-gateway/lastcluster/local_store.go b/src/invocation-plane-services/llm-api-gateway/lastcluster/local_store.go new file mode 100644 index 0000000000..bb5611ec79 --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/lastcluster/local_store.go @@ -0,0 +1,112 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package lastcluster + +import ( + "container/list" + "context" + "sync" + "time" +) + +// LocalStore is an in-process LRU with per-entry TTL, used when Olric is off. +// It holds at most maxEntries entries and evicts the least recently used one +// first. Reads refresh recency but not the TTL; only writes reset the TTL. +type LocalStore struct { + mu sync.Mutex + maxEntries int + now func() time.Time + order *list.List // front is most recently used + entries map[string]*list.Element +} + +type localEntry struct { + key string + value string + expiresAt time.Time +} + +// LocalOption configures a LocalStore. +type LocalOption func(*LocalStore) + +// WithClock overrides the clock used for TTL checks. Tests only. +func WithClock(now func() time.Time) LocalOption { + return func(s *LocalStore) { s.now = now } +} + +// NewLocalStore returns an empty LocalStore. maxEntries below 1 is treated as 1. +func NewLocalStore(maxEntries int, opts ...LocalOption) *LocalStore { + if maxEntries < 1 { + maxEntries = 1 + } + s := &LocalStore{ + maxEntries: maxEntries, + now: time.Now, + order: list.New(), + entries: make(map[string]*list.Element), + } + for _, opt := range opts { + opt(s) + } + return s +} + +func (s *LocalStore) Get(_ context.Context, key string) (string, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + elem, ok := s.entries[key] + if !ok { + return "", false, nil + } + entry := elem.Value.(*localEntry) + if !s.now().Before(entry.expiresAt) { + s.order.Remove(elem) + delete(s.entries, key) + return "", false, nil + } + s.order.MoveToFront(elem) + return entry.value, true, nil +} + +func (s *LocalStore) Set(_ context.Context, key, value string, ttl time.Duration) error { + s.mu.Lock() + defer s.mu.Unlock() + expiresAt := s.now().Add(ttl) + if elem, ok := s.entries[key]; ok { + entry := elem.Value.(*localEntry) + entry.value = value + entry.expiresAt = expiresAt + s.order.MoveToFront(elem) + return nil + } + s.entries[key] = s.order.PushFront(&localEntry{key: key, value: value, expiresAt: expiresAt}) + for s.order.Len() > s.maxEntries { + oldest := s.order.Back() + s.order.Remove(oldest) + delete(s.entries, oldest.Value.(*localEntry).key) + } + return nil +} + +// Len returns the number of entries held, including expired entries that +// have not been read or evicted yet. +func (s *LocalStore) Len() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.order.Len() +} diff --git a/src/invocation-plane-services/llm-api-gateway/lastcluster/local_store_test.go b/src/invocation-plane-services/llm-api-gateway/lastcluster/local_store_test.go new file mode 100644 index 0000000000..35d9b89126 --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/lastcluster/local_store_test.go @@ -0,0 +1,149 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package lastcluster + +import ( + "context" + "fmt" + "strings" + "testing" + "time" +) + +func TestKeyIsPrefixedHashOfLengthPrefixedParts(t *testing.T) { + key := Key("routing-key", "model", "mt:v1:session:abc") + if !strings.HasPrefix(key, "lc:v1:") || len(key) != len("lc:v1:")+64 { + t.Fatalf("Key() = %q, want lc:v1: plus 64 hex characters", key) + } + if strings.Contains(key, "routing-key") { + t.Fatalf("Key() = %q leaks the raw routing key", key) + } + if Key("ab", "c", "d") == Key("a", "bc", "d") { + t.Fatal("Key() collides for different splits of the same bytes") + } + if Key("r", "m1", "k") == Key("r", "m2", "k") || Key("r1", "m", "k") == Key("r2", "m", "k") { + t.Fatal("Key() must differ by model and routing key") + } + if Key("r", "m", "k") != Key("r", "m", "k") { + t.Fatal("Key() must be deterministic") + } +} + +func TestValidClusterID(t *testing.T) { + cases := map[string]bool{ + "": false, + "cluster-b": true, + "cluster b": true, + strings.Repeat("a", MaxValueBytes): true, + strings.Repeat("a", MaxValueBytes+1): false, + "bad\nvalue": false, + "caf\xc3\xa9": false, + } + for value, want := range cases { + if got := validClusterID(value); got != want { + t.Errorf("validClusterID(%q) = %v, want %v", value, got, want) + } + } +} + +type fakeClock struct{ now time.Time } + +func (c *fakeClock) Now() time.Time { return c.now } +func (c *fakeClock) Advance(d time.Duration) { c.now = c.now.Add(d) } + +func TestLocalStoreTTL(t *testing.T) { + ctx := context.Background() + clock := &fakeClock{now: time.Unix(1000, 0)} + store := NewLocalStore(10, WithClock(clock.Now)) + ttl := 10 * time.Minute + + if err := store.Set(ctx, "k", "B", ttl); err != nil { + t.Fatalf("Set() error = %v", err) + } + clock.Advance(ttl - time.Nanosecond) + if value, ok, _ := store.Get(ctx, "k"); !ok || value != "B" { + t.Fatalf("Get() before TTL = %q, %v; want B, true", value, ok) + } + clock.Advance(time.Nanosecond) + if value, ok, _ := store.Get(ctx, "k"); ok { + t.Fatalf("Get() at TTL = %q, true; want miss", value) + } + if store.Len() != 0 { + t.Fatalf("Len() = %d after expired read, want 0", store.Len()) + } +} + +func TestLocalStoreWriteResetsTTL(t *testing.T) { + ctx := context.Background() + clock := &fakeClock{now: time.Unix(1000, 0)} + store := NewLocalStore(10, WithClock(clock.Now)) + ttl := time.Minute + + _ = store.Set(ctx, "k", "A", ttl) + clock.Advance(ttl / 2) + _ = store.Set(ctx, "k", "B", ttl) + clock.Advance(ttl/2 + ttl/4) + if value, ok, _ := store.Get(ctx, "k"); !ok || value != "B" { + t.Fatalf("Get() after rewrite = %q, %v; want B, true", value, ok) + } + // A read does not extend the TTL. + clock.Advance(ttl / 4) + if _, ok, _ := store.Get(ctx, "k"); ok { + t.Fatal("Get() after rewrite TTL = hit, want miss") + } +} + +func TestLocalStoreEvictsLeastRecentlyUsed(t *testing.T) { + ctx := context.Background() + store := NewLocalStore(2) + _ = store.Set(ctx, "a", "A", time.Minute) + _ = store.Set(ctx, "b", "B", time.Minute) + if _, ok, _ := store.Get(ctx, "a"); !ok { + t.Fatal("Get(a) = miss, want hit") + } + _ = store.Set(ctx, "c", "C", time.Minute) + + if _, ok, _ := store.Get(ctx, "b"); ok { + t.Fatal("Get(b) = hit, want b evicted as least recently used") + } + for _, key := range []string{"a", "c"} { + if _, ok, _ := store.Get(ctx, key); !ok { + t.Fatalf("Get(%s) = miss, want hit", key) + } + } +} + +func TestLocalStoreNeverExceedsCap(t *testing.T) { + ctx := context.Background() + store := NewLocalStore(100) + for i := 0; i < 1000; i++ { + _ = store.Set(ctx, fmt.Sprintf("k%d", i), "B", time.Minute) + if store.Len() > 100 { + t.Fatalf("Len() = %d after %d writes, want <= 100", store.Len(), i+1) + } + } + if store.Len() != 100 { + t.Fatalf("Len() = %d, want 100", store.Len()) + } + if _, ok, _ := store.Get(ctx, "k999"); !ok { + t.Fatal("newest entry evicted") + } + if _, ok, _ := store.Get(ctx, "k0"); ok { + t.Fatal("oldest entry still present") + } +} diff --git a/src/invocation-plane-services/llm-api-gateway/lastcluster/olric_store.go b/src/invocation-plane-services/llm-api-gateway/lastcluster/olric_store.go new file mode 100644 index 0000000000..522ddfc2db --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/lastcluster/olric_store.go @@ -0,0 +1,59 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package lastcluster + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/olric-data/olric" +) + +// OlricStore keeps entries in an Olric DMap shared by every gateway replica +// in the Olric cluster. +type OlricStore struct { + dmap olric.DMap +} + +func NewOlricStore(dmap olric.DMap) *OlricStore { + return &OlricStore{dmap: dmap} +} + +func (s *OlricStore) Get(ctx context.Context, key string) (string, bool, error) { + resp, err := s.dmap.Get(ctx, key) + if err != nil { + if errors.Is(err, olric.ErrKeyNotFound) { + return "", false, nil + } + return "", false, fmt.Errorf("olric get: %w", err) + } + value, err := resp.String() + if err != nil { + return "", false, fmt.Errorf("olric get decode: %w", err) + } + return value, true, nil +} + +func (s *OlricStore) Set(ctx context.Context, key, value string, ttl time.Duration) error { + if err := s.dmap.Put(ctx, key, value, olric.PX(ttl)); err != nil { + return fmt.Errorf("olric put: %w", err) + } + return nil +} diff --git a/src/invocation-plane-services/llm-api-gateway/lastcluster/olric_store_test.go b/src/invocation-plane-services/llm-api-gateway/lastcluster/olric_store_test.go new file mode 100644 index 0000000000..272fc7e963 --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/lastcluster/olric_store_test.go @@ -0,0 +1,93 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package lastcluster + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/internal/olrictest" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/requestctx" +) + +func TestOlricStoreTTL(t *testing.T) { + if testing.Short() { + t.Skip("skipping embedded Olric test in -short mode") + } + ctx := context.Background() + node := olrictest.StartCluster(t, 1)[0] + dm, err := node.Client.NewDMap(DMapName) + if err != nil { + t.Fatalf("NewDMap() error = %v", err) + } + store := NewOlricStore(dm) + + if _, ok, err := store.Get(ctx, "missing"); ok || err != nil { + t.Fatalf("Get(missing) = %v, %v; want miss without error", ok, err) + } + if err := store.Set(ctx, "k", "cluster-b", 300*time.Millisecond); err != nil { + t.Fatalf("Set() error = %v", err) + } + if value, ok, err := store.Get(ctx, "k"); err != nil || !ok || value != "cluster-b" { + t.Fatalf("Get() = %q, %v, %v; want cluster-b", value, ok, err) + } + deadline := time.Now().Add(5 * time.Second) + for { + _, ok, err := store.Get(ctx, "k") + if err != nil { + t.Fatalf("Get() error = %v", err) + } + if !ok { + return + } + if time.Now().After(deadline) { + t.Fatal("entry did not expire after its TTL") + } + time.Sleep(50 * time.Millisecond) + } +} + +func TestTwoReplicasShareLastCluster(t *testing.T) { + if testing.Short() { + t.Skip("skipping multi-node Olric test in -short mode") + } + ctx := context.Background() + nodes := olrictest.StartCluster(t, 2) + trackers := make([]*Tracker, len(nodes)) + for i, node := range nodes { + dm, err := node.Client.NewDMap(DMapName) + if err != nil { + t.Fatalf("NewDMap() on node %d error = %v", i, err) + } + trackers[i] = NewTracker(NewOlricStore(dm), Options{TTL: time.Minute, LookupTimeout: time.Second}) + } + reqCtx := sessionCtx(requestctx.SessionSourcePromptCacheKey) + + key, clusterID := trackers[0].Lookup(ctx, reqCtx, reqCtx.Model) + if clusterID != "" { + t.Fatalf("first lookup cluster = %q, want miss", clusterID) + } + trackers[0].Remember(ctx, reqCtx, reqCtx.Model, key, http.StatusOK, "cluster-b") + trackers[0].Drain(ctx) + + if _, clusterID := trackers[1].Lookup(ctx, reqCtx, reqCtx.Model); clusterID != "cluster-b" { + t.Fatalf("replica 2 lookup cluster = %q, want cluster-b", clusterID) + } +} diff --git a/src/invocation-plane-services/llm-api-gateway/lastcluster/store.go b/src/invocation-plane-services/llm-api-gateway/lastcluster/store.go new file mode 100644 index 0000000000..fb344685a0 --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/lastcluster/store.go @@ -0,0 +1,77 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package lastcluster remembers which Stargate cluster served each session's +// last successful request so the gateway can send it back to Stargate as the +// x-stargate-last-cluster-id hint. +package lastcluster + +import ( + "context" + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "io" + "time" +) + +const ( + // DMapName is the Olric DMap that holds last-cluster entries. It is + // separate from the rate-limit DMap. + DMapName = "stargate-last-cluster" + // MaxValueBytes bounds the stored cluster ID. Longer values are not stored. + MaxValueBytes = 256 + + keyPrefix = "lc:v1:" +) + +// Store holds session key to cluster ID entries with a per-entry TTL. +type Store interface { + // Get returns ok=false when the key is absent or expired. + Get(ctx context.Context, key string) (value string, ok bool, err error) + // Set writes the value and resets the entry TTL. + Set(ctx context.Context, key, value string, ttl time.Duration) error +} + +// Key derives the store key for a session. Each part is length-prefixed +// before hashing so different splits of the same bytes never collide, and raw +// routing keys or session identifiers never appear in the store key. +func Key(routingKey, model, cacheAffinityKey string) string { + h := sha256.New() + var n [8]byte + for _, part := range []string{routingKey, model, cacheAffinityKey} { + binary.BigEndian.PutUint64(n[:], uint64(len(part))) + _, _ = h.Write(n[:]) + _, _ = io.WriteString(h, part) + } + return keyPrefix + hex.EncodeToString(h.Sum(nil)) +} + +// validClusterID reports whether a cluster ID may be stored and later sent as +// a request header value: non-empty, at most MaxValueBytes, and printable +// ASCII only. +func validClusterID(value string) bool { + if value == "" || len(value) > MaxValueBytes { + return false + } + for i := 0; i < len(value); i++ { + if c := value[i]; c < 0x20 || c > 0x7e { + return false + } + } + return true +} diff --git a/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker.go b/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker.go new file mode 100644 index 0000000000..b19c966563 --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker.go @@ -0,0 +1,267 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package lastcluster + +import ( + "context" + "errors" + "strings" + "sync" + "time" + + "github.com/rs/zerolog" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" + + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/requestctx" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/telemetry" +) + +// DefaultWriteTimeout bounds one background store write. +const DefaultWriteTimeout = time.Second + +// Options configures a Tracker. +type Options struct { + TTL time.Duration + LookupTimeout time.Duration + // WriteTimeout bounds one background write. Zero uses DefaultWriteTimeout. + WriteTimeout time.Duration +} + +// Tracker looks up and records the last cluster per session. A nil Tracker is +// valid and does nothing. +type Tracker struct { + store Store + opts Options + + // inflight counts background writes. idle is closed when it drops to + // zero. A WaitGroup is not used because writes may start while Drain is + // waiting during shutdown. + mu sync.Mutex + inflight int + idle chan struct{} +} + +func NewTracker(store Store, opts Options) *Tracker { + if opts.WriteTimeout <= 0 { + opts.WriteTimeout = DefaultWriteTimeout + } + return &Tracker{store: store, opts: opts} +} + +// eligibleSource reports whether a session source yields a key that stays +// stable across turns. Payload-derived keys hash the full message list, which +// changes every turn, so they never return. +func eligibleSource(source string) bool { + switch source { + case requestctx.SessionSourcePromptCacheKey, + requestctx.SessionSourceConversationID, + requestctx.SessionSourceHeader: + return true + default: + return false + } +} + +type lookupResult struct { + value string + ok bool + err error +} + +// Lookup returns the store key for the session and the remembered cluster ID. +// key is empty when the request has no eligible session, in which case no +// write should follow. clusterID is empty on a miss, error, or timeout. The +// lookup never takes longer than the configured lookup timeout. +func (t *Tracker) Lookup( + ctx context.Context, + reqCtx *requestctx.RequestContext, + model string, +) (key string, clusterID string) { + if t == nil || reqCtx == nil || reqCtx.CacheAffinityKey == "" { + return "", "" + } + if !eligibleSource(reqCtx.SessionSource) || model == "" { + telemetry.RecordLastClusterLookup(ctx, telemetry.LastClusterLookupSkipped, 0) + return "", "" + } + // The deadline starts before any span or key work so the whole lookup, + // not only the store call, fits in the lookup timeout. + start := time.Now() + lookupCtx, cancel := context.WithDeadline(ctx, start.Add(t.opts.LookupTimeout)) + defer cancel() + key = Key(reqCtx.RoutingKey, model, reqCtx.CacheAffinityKey) + + ctx, span := telemetry.Tracer().Start(ctx, "llm-api-gateway.last_cluster_lookup") + defer span.End() + setFunctionID(span, reqCtx.RoutingKey) + lookupCtx = trace.ContextWithSpan(lookupCtx, span) + + // The store call runs in its own goroutine so a store that ignores + // context cancellation still cannot hold the request past the deadline. + done := make(chan lookupResult, 1) + go func() { + value, ok, err := t.store.Get(lookupCtx, key) + done <- lookupResult{value: value, ok: ok, err: err} + }() + + var res lookupResult + select { + case res = <-done: + case <-lookupCtx.Done(): + // Prefer a result that landed at the deadline. + select { + case res = <-done: + default: + res.err = lookupCtx.Err() + } + } + + result := telemetry.LastClusterLookupMiss + switch { + case errors.Is(res.err, context.DeadlineExceeded), errors.Is(res.err, context.Canceled): + result = telemetry.LastClusterLookupTimeout + case res.err != nil: + result = telemetry.LastClusterLookupError + case res.ok && validClusterID(res.value): + result = telemetry.LastClusterLookupHit + clusterID = res.value + } + telemetry.RecordLastClusterLookup(ctx, result, time.Since(start)) + span.SetAttributes(attribute.String("result", result)) + + // A cancelled request (client gone) is not a store failure. + if res.err == nil || errors.Is(ctx.Err(), context.Canceled) { + return key, clusterID + } + span.RecordError(res.err) + span.SetStatus(codes.Error, "last cluster lookup failed") + logger := telemetry.Logger(ctx) + if result == telemetry.LastClusterLookupTimeout { + // A slow store times out every lookup at once; the counter and span + // already count each one, so sample the log. + sampled := logger.Sample(zerolog.Often) + logger = &sampled + } + logger.Warn(). + Err(res.err). + Str("request_id", reqCtx.RequestID). + Str("model", model). + Str("routing_key", reqCtx.RoutingKey). + Str("result", result). + Msg("stargate last-cluster lookup failed") + return key, clusterID +} + +// setFunctionID tags the span with the function ID, which the gateway takes +// from the routing key. +func setFunctionID(span trace.Span, routingKey string) { + if routingKey != "" { + span.SetAttributes(attribute.String("nvcf.function.id", routingKey)) + } +} + +// Remember records clusterID for the session in the background when status is +// 2xx. key and model are the values Lookup used. It never blocks the caller. +// Empty, oversized, or non-printable cluster IDs are ignored. +func (t *Tracker) Remember( + ctx context.Context, + reqCtx *requestctx.RequestContext, + model string, + key string, + status int, + clusterID string, +) { + if t == nil || key == "" || status < 200 || status >= 300 { + return + } + clusterID = strings.TrimSpace(clusterID) + if !validClusterID(clusterID) { + return + } + var requestID, routingKey string + if reqCtx != nil { + requestID, routingKey = reqCtx.RequestID, reqCtx.RoutingKey + } + + // The write outlives the request: keep the trace and logger from ctx but + // not its cancellation. + writeCtx := context.WithoutCancel(ctx) + t.beginWrite() + go func() { + defer t.endWrite() + writeCtx, cancel := context.WithTimeout(writeCtx, t.opts.WriteTimeout) + defer cancel() + writeCtx, span := telemetry.Tracer().Start(writeCtx, "llm-api-gateway.last_cluster_write") + defer span.End() + setFunctionID(span, routingKey) + + result := telemetry.LastClusterWriteOK + err := t.store.Set(writeCtx, key, clusterID, t.opts.TTL) + if err != nil { + result = telemetry.LastClusterWriteError + span.RecordError(err) + span.SetStatus(codes.Error, "last cluster write failed") + telemetry.Logger(writeCtx).Warn(). + Err(err). + Str("request_id", requestID). + Str("model", model). + Str("routing_key", routingKey). + Msg("stargate last-cluster write failed") + } + span.SetAttributes(attribute.String("result", result)) + telemetry.RecordLastClusterWrite(writeCtx, result) + }() +} + +func (t *Tracker) beginWrite() { + t.mu.Lock() + defer t.mu.Unlock() + if t.inflight == 0 { + t.idle = make(chan struct{}) + } + t.inflight++ +} + +func (t *Tracker) endWrite() { + t.mu.Lock() + defer t.mu.Unlock() + t.inflight-- + if t.inflight == 0 { + close(t.idle) + } +} + +// Drain waits until no background write is in flight or ctx is done. +func (t *Tracker) Drain(ctx context.Context) { + if t == nil { + return + } + t.mu.Lock() + if t.inflight == 0 { + t.mu.Unlock() + return + } + idle := t.idle + t.mu.Unlock() + select { + case <-idle: + case <-ctx.Done(): + } +} diff --git a/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker_test.go b/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker_test.go new file mode 100644 index 0000000000..79f4c426cf --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker_test.go @@ -0,0 +1,523 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package lastcluster + +import ( + "bytes" + "context" + "errors" + "net/http" + "strings" + "sync" + "testing" + "time" + + "github.com/rs/zerolog" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + "go.opentelemetry.io/otel/trace" + + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/requestctx" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/telemetry" +) + +// fakeStore records calls and can inject errors or block Get while ignoring +// the context, like a store that does not honor cancellation. +type fakeStore struct { + mu sync.Mutex + values map[string]string + gets int + sets int + getErr error + setErr error + blockGet chan struct{} +} + +func newFakeStore() *fakeStore { + return &fakeStore{values: map[string]string{}} +} + +func (s *fakeStore) Get(_ context.Context, key string) (string, bool, error) { + s.mu.Lock() + s.gets++ + block, err := s.blockGet, s.getErr + value, ok := s.values[key] + s.mu.Unlock() + if block != nil { + <-block + } + if err != nil { + return "", false, err + } + return value, ok, nil +} + +func (s *fakeStore) Set(_ context.Context, key, value string, _ time.Duration) error { + s.mu.Lock() + defer s.mu.Unlock() + s.sets++ + if s.setErr != nil { + return s.setErr + } + s.values[key] = value + return nil +} + +func (s *fakeStore) counts() (gets, sets int) { + s.mu.Lock() + defer s.mu.Unlock() + return s.gets, s.sets +} + +func (s *fakeStore) value(key string) string { + s.mu.Lock() + defer s.mu.Unlock() + return s.values[key] +} + +func sessionCtx(source string) *requestctx.RequestContext { + return &requestctx.RequestContext{ + RequestID: "req-1", + RoutingKey: "fn-alpha", + Model: "company/model", + SessionID: "raw-session-secret", + SessionSource: source, + CacheAffinityKey: "mt:v1:session:affinity-secret", + } +} + +func newTestTracker(store Store) *Tracker { + return NewTracker(store, Options{TTL: time.Minute, LookupTimeout: 20 * time.Millisecond}) +} + +func installSpanRecorder(t *testing.T) *tracetest.SpanRecorder { + t.Helper() + recorder := tracetest.NewSpanRecorder() + provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder)) + oldTracer := telemetry.Tracer + telemetry.Tracer = sync.OnceValue(func() trace.Tracer { return provider.Tracer("test") }) + t.Cleanup(func() { + telemetry.Tracer = oldTracer + _ = provider.Shutdown(context.Background()) + }) + return recorder +} + +func installMetricReader(t *testing.T) *sdkmetric.ManualReader { + t.Helper() + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + oldProvider := otel.GetMeterProvider() + otel.SetMeterProvider(provider) + t.Cleanup(func() { + otel.SetMeterProvider(oldProvider) + _ = provider.Shutdown(context.Background()) + }) + return reader +} + +func counterValues(t *testing.T, reader *sdkmetric.ManualReader, name string) map[string]int64 { + t.Helper() + var rm metricdata.ResourceMetrics + if err := reader.Collect(context.Background(), &rm); err != nil { + t.Fatalf("Collect() error = %v", err) + } + out := map[string]int64{} + for _, sm := range rm.ScopeMetrics { + for _, m := range sm.Metrics { + if m.Name != name { + continue + } + for _, dp := range m.Data.(metricdata.Sum[int64]).DataPoints { + value, _ := dp.Attributes.Value("result") + out[value.AsString()] = dp.Value + } + } + } + return out +} + +func spanByName(t *testing.T, recorder *tracetest.SpanRecorder, name string) sdktrace.ReadOnlySpan { + t.Helper() + for _, span := range recorder.Ended() { + if span.Name() == name { + return span + } + } + t.Fatalf("span %q not recorded", name) + return nil +} + +func spanResult(span sdktrace.ReadOnlySpan) string { + for _, attr := range span.Attributes() { + if attr.Key == attribute.Key("result") { + return attr.Value.AsString() + } + } + return "" +} + +func TestTrackerRoundTripForEligibleSources(t *testing.T) { + for _, source := range []string{ + requestctx.SessionSourcePromptCacheKey, + requestctx.SessionSourceConversationID, + requestctx.SessionSourceHeader, + } { + t.Run(source, func(t *testing.T) { + ctx := context.Background() + tracker := newTestTracker(NewLocalStore(10)) + reqCtx := sessionCtx(source) + + key, clusterID := tracker.Lookup(ctx, reqCtx, reqCtx.Model) + if key == "" || clusterID != "" { + t.Fatalf("first Lookup() = %q, %q; want key and no cluster", key, clusterID) + } + tracker.Remember(ctx, reqCtx, reqCtx.Model, key, http.StatusOK, "cluster-b") + tracker.Drain(ctx) + + _, clusterID = tracker.Lookup(ctx, reqCtx, reqCtx.Model) + if clusterID != "cluster-b" { + t.Fatalf("second Lookup() cluster = %q, want cluster-b", clusterID) + } + }) + } +} + +func TestTrackerSkipsIneligibleSessions(t *testing.T) { + reader := installMetricReader(t) + ctx := context.Background() + store := newFakeStore() + tracker := newTestTracker(store) + + for _, source := range []string{ + requestctx.SessionSourcePayload, + requestctx.SessionSourceClaudeCodeHeader, + "unknown", + } { + key, clusterID := tracker.Lookup(ctx, sessionCtx(source), "company/model") + if key != "" || clusterID != "" { + t.Fatalf("Lookup(%s) = %q, %q; want skipped", source, key, clusterID) + } + tracker.Remember(ctx, sessionCtx(source), "company/model", key, http.StatusOK, "cluster-b") + } + if key, _ := tracker.Lookup(ctx, sessionCtx(requestctx.SessionSourceHeader), ""); key != "" { + t.Fatal("Lookup() with empty model must be skipped") + } + // No session at all is not counted. + if key, _ := tracker.Lookup(ctx, &requestctx.RequestContext{RoutingKey: "fn"}, "m"); key != "" { + t.Fatal("Lookup() without a session must return no key") + } + tracker.Drain(ctx) + + if gets, sets := store.counts(); gets != 0 || sets != 0 { + t.Fatalf("store gets=%d sets=%d, want 0 and 0", gets, sets) + } + lookups := counterValues(t, reader, "llm_api_gateway_last_cluster_lookups_total") + if lookups["skipped"] != 4 || len(lookups) != 1 { + t.Fatalf("lookups = %v, want only skipped=4", lookups) + } +} + +func TestTrackerNilIsNoop(t *testing.T) { + var tracker *Tracker + ctx := context.Background() + if key, clusterID := tracker.Lookup(ctx, sessionCtx(requestctx.SessionSourceHeader), "m"); key != "" || clusterID != "" { + t.Fatal("nil Tracker Lookup() must return empty values") + } + tracker.Remember(ctx, nil, "m", "k", http.StatusOK, "B") + tracker.Drain(ctx) +} + +func TestTrackerRememberOnlyStoresValidSuccessfulResponses(t *testing.T) { + ctx := context.Background() + store := newFakeStore() + tracker := newTestTracker(store) + reqCtx := sessionCtx(requestctx.SessionSourceHeader) + key, _ := tracker.Lookup(ctx, reqCtx, reqCtx.Model) + + tracker.Remember(ctx, reqCtx, reqCtx.Model, key, http.StatusOK, "cluster-a") + tracker.Drain(ctx) + for _, tc := range []struct { + status int + clusterID string + }{ + {status: http.StatusServiceUnavailable, clusterID: "cluster-c"}, + {status: http.StatusTooManyRequests, clusterID: "cluster-c"}, + {status: http.StatusMultipleChoices, clusterID: "cluster-c"}, + {status: http.StatusOK, clusterID: ""}, + {status: http.StatusOK, clusterID: " "}, + {status: http.StatusOK, clusterID: strings.Repeat("c", MaxValueBytes+1)}, + {status: http.StatusOK, clusterID: "bad\x01value"}, + } { + tracker.Remember(ctx, reqCtx, reqCtx.Model, key, tc.status, tc.clusterID) + } + tracker.Drain(ctx) + + if got := store.value(key); got != "cluster-a" { + t.Fatalf("stored value = %q, want cluster-a", got) + } + if _, sets := store.counts(); sets != 1 { + t.Fatalf("store sets = %d, want 1", sets) + } + + tracker.Remember(ctx, reqCtx, reqCtx.Model, key, http.StatusCreated, " cluster-b ") + tracker.Drain(ctx) + if got := store.value(key); got != "cluster-b" { + t.Fatalf("stored value = %q, want trimmed cluster-b", got) + } +} + +func TestTrackerTreatsInvalidStoredValueAsMiss(t *testing.T) { + ctx := context.Background() + store := newFakeStore() + tracker := newTestTracker(store) + reqCtx := sessionCtx(requestctx.SessionSourceHeader) + key := Key(reqCtx.RoutingKey, reqCtx.Model, reqCtx.CacheAffinityKey) + store.values[key] = "bad\nvalue" + + if _, clusterID := tracker.Lookup(ctx, reqCtx, reqCtx.Model); clusterID != "" { + t.Fatalf("Lookup() cluster = %q, want miss for an invalid stored value", clusterID) + } +} + +func TestLookupTimeout(t *testing.T) { + recorder := installSpanRecorder(t) + reader := installMetricReader(t) + store := newFakeStore() + store.blockGet = make(chan struct{}) + t.Cleanup(func() { close(store.blockGet) }) + tracker := newTestTracker(store) + reqCtx := sessionCtx(requestctx.SessionSourcePromptCacheKey) + + start := time.Now() + key, clusterID := tracker.Lookup(context.Background(), reqCtx, reqCtx.Model) + elapsed := time.Since(start) + + if elapsed >= 25*time.Millisecond { + t.Fatalf("Lookup() took %s, want < timeout + 5ms", elapsed) + } + if key == "" || clusterID != "" { + t.Fatalf("Lookup() = %q, %q; want key and no cluster", key, clusterID) + } + span := spanByName(t, recorder, "llm-api-gateway.last_cluster_lookup") + if spanResult(span) != "timeout" || span.Status().Code != codes.Error { + t.Fatalf("span result=%q status=%v, want timeout and Error", spanResult(span), span.Status().Code) + } + if got := counterValues(t, reader, "llm_api_gateway_last_cluster_lookups_total"); got["timeout"] != 1 { + t.Fatalf("lookups = %v, want timeout=1", got) + } +} + +func TestLookupCancelledRequestIsNotAStoreFailure(t *testing.T) { + recorder := installSpanRecorder(t) + reader := installMetricReader(t) + var buf syncBuffer + store := newFakeStore() + store.blockGet = make(chan struct{}) + t.Cleanup(func() { close(store.blockGet) }) + tracker := newTestTracker(store) + reqCtx := sessionCtx(requestctx.SessionSourceHeader) + ctx, cancel := context.WithCancel(zerolog.New(&buf).WithContext(context.Background())) + time.AfterFunc(time.Millisecond, cancel) + + key, clusterID := tracker.Lookup(ctx, reqCtx, reqCtx.Model) + if key == "" || clusterID != "" { + t.Fatalf("Lookup() = %q, %q; want key and no cluster", key, clusterID) + } + span := spanByName(t, recorder, "llm-api-gateway.last_cluster_lookup") + if spanResult(span) != "timeout" || span.Status().Code == codes.Error { + t.Fatalf("span result=%q status=%v, want timeout without error status", spanResult(span), span.Status().Code) + } + if got := counterValues(t, reader, "llm_api_gateway_last_cluster_lookups_total"); got["timeout"] != 1 || got["error"] != 0 { + t.Fatalf("lookups = %v, want timeout=1 and no error", got) + } + if logs := buf.String(); strings.Contains(logs, `"level":"warn"`) { + t.Fatalf("logs = %s, want no warn for a cancelled request", logs) + } +} + +func TestTrackerSpansCarryFunctionID(t *testing.T) { + recorder := installSpanRecorder(t) + ctx := context.Background() + tracker := newTestTracker(newFakeStore()) + reqCtx := sessionCtx(requestctx.SessionSourceHeader) + + key, _ := tracker.Lookup(ctx, reqCtx, reqCtx.Model) + tracker.Remember(ctx, reqCtx, reqCtx.Model, key, http.StatusOK, "cluster-b") + tracker.Drain(ctx) + + for _, name := range []string{"llm-api-gateway.last_cluster_lookup", "llm-api-gateway.last_cluster_write"} { + span := spanByName(t, recorder, name) + found := false + for _, attr := range span.Attributes() { + if attr.Key == "nvcf.function.id" && attr.Value.AsString() == reqCtx.RoutingKey { + found = true + } + } + if !found { + t.Fatalf("span %s attributes = %v, want nvcf.function.id=%s", name, span.Attributes(), reqCtx.RoutingKey) + } + } +} + +func TestTrackerWriteLogUsesLookupModel(t *testing.T) { + var buf syncBuffer + ctx := zerolog.New(&buf).WithContext(context.Background()) + store := newFakeStore() + store.setErr = errors.New("olric unavailable") + tracker := newTestTracker(store) + reqCtx := sessionCtx(requestctx.SessionSourceHeader) + + key, _ := tracker.Lookup(ctx, reqCtx, "company/effective-model") + tracker.Remember(ctx, reqCtx, "company/effective-model", key, http.StatusOK, "cluster-b") + tracker.Drain(ctx) + + if logs := buf.String(); !strings.Contains(logs, `"model":"company/effective-model"`) { + t.Fatalf("logs = %s, want the model used for the key", logs) + } +} + +func TestTrackerSpansAndMetrics(t *testing.T) { + recorder := installSpanRecorder(t) + reader := installMetricReader(t) + ctx := context.Background() + store := newFakeStore() + tracker := newTestTracker(store) + reqCtx := sessionCtx(requestctx.SessionSourceHeader) + + key, _ := tracker.Lookup(ctx, reqCtx, reqCtx.Model) // miss + tracker.Remember(ctx, reqCtx, reqCtx.Model, key, http.StatusOK, "cluster-b") + tracker.Drain(ctx) + tracker.Lookup(ctx, reqCtx, reqCtx.Model) // hit + + store.mu.Lock() + store.getErr = errors.New("olric unavailable") + store.setErr = errors.New("olric unavailable") + store.mu.Unlock() + tracker.Lookup(ctx, reqCtx, reqCtx.Model) // error + tracker.Remember(ctx, reqCtx, reqCtx.Model, key, http.StatusOK, "cluster-c") + tracker.Drain(ctx) + + var lookups, writes []sdktrace.ReadOnlySpan + for _, span := range recorder.Ended() { + switch span.Name() { + case "llm-api-gateway.last_cluster_lookup": + lookups = append(lookups, span) + case "llm-api-gateway.last_cluster_write": + writes = append(writes, span) + } + } + if len(lookups) != 3 || len(writes) != 2 { + t.Fatalf("spans lookups=%d writes=%d, want 3 and 2", len(lookups), len(writes)) + } + for i, want := range []struct { + result string + code codes.Code + }{{"miss", codes.Unset}, {"hit", codes.Unset}, {"error", codes.Error}} { + if got := spanResult(lookups[i]); got != want.result || lookups[i].Status().Code != want.code { + t.Fatalf("lookup span %d result=%q status=%v, want %q %v", i, got, lookups[i].Status().Code, want.result, want.code) + } + } + for i, want := range []struct { + result string + code codes.Code + }{{"ok", codes.Unset}, {"error", codes.Error}} { + if got := spanResult(writes[i]); got != want.result || writes[i].Status().Code != want.code { + t.Fatalf("write span %d result=%q status=%v, want %q %v", i, got, writes[i].Status().Code, want.result, want.code) + } + } + + gotLookups := counterValues(t, reader, "llm_api_gateway_last_cluster_lookups_total") + if gotLookups["miss"] != 1 || gotLookups["hit"] != 1 || gotLookups["error"] != 1 { + t.Fatalf("lookups = %v, want miss=1 hit=1 error=1", gotLookups) + } + gotWrites := counterValues(t, reader, "llm_api_gateway_last_cluster_writes_total") + if gotWrites["ok"] != 1 || gotWrites["error"] != 1 { + t.Fatalf("writes = %v, want ok=1 error=1", gotWrites) + } +} + +func TestTrackerStoreErrorsLogWarnWithoutSessionIdentifiers(t *testing.T) { + var buf syncBuffer + ctx := zerolog.New(&buf).WithContext(context.Background()) + store := newFakeStore() + store.getErr = errors.New("olric unavailable") + store.setErr = errors.New("olric unavailable") + tracker := newTestTracker(store) + reqCtx := sessionCtx(requestctx.SessionSourceHeader) + + key, clusterID := tracker.Lookup(ctx, reqCtx, reqCtx.Model) + if key == "" || clusterID != "" { + t.Fatalf("Lookup() = %q, %q; want key and no cluster on error", key, clusterID) + } + tracker.Remember(ctx, reqCtx, reqCtx.Model, key, http.StatusOK, "cluster-b") + tracker.Drain(ctx) + + logs := buf.String() + if strings.Count(logs, `"level":"warn"`) != 2 { + t.Fatalf("logs = %s, want two warn lines", logs) + } + for _, want := range []string{`"request_id":"req-1"`, `"model":"company/model"`, `"routing_key":"fn-alpha"`} { + if strings.Count(logs, want) != 2 { + t.Fatalf("logs = %s, want %s on both lines", logs, want) + } + } + for _, secret := range []string{reqCtx.CacheAffinityKey, reqCtx.SessionID, "affinity-secret", key} { + if strings.Contains(logs, secret) { + t.Fatalf("logs leak %q: %s", secret, logs) + } + } +} + +func TestTrackerWriteSurvivesRequestCancellation(t *testing.T) { + store := newFakeStore() + tracker := newTestTracker(store) + reqCtx := sessionCtx(requestctx.SessionSourceHeader) + ctx, cancel := context.WithCancel(context.Background()) + key, _ := tracker.Lookup(ctx, reqCtx, reqCtx.Model) + + tracker.Remember(ctx, reqCtx, reqCtx.Model, key, http.StatusOK, "cluster-b") + cancel() + tracker.Drain(context.Background()) + + if got := store.value(key); got != "cluster-b" { + t.Fatalf("stored value = %q, want cluster-b", got) + } +} + +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *syncBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *syncBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} diff --git a/src/invocation-plane-services/llm-api-gateway/provider/BUILD.bazel b/src/invocation-plane-services/llm-api-gateway/provider/BUILD.bazel index 7014a8b150..a3709a74e2 100644 --- a/src/invocation-plane-services/llm-api-gateway/provider/BUILD.bazel +++ b/src/invocation-plane-services/llm-api-gateway/provider/BUILD.bazel @@ -28,6 +28,7 @@ go_library( "//src/invocation-plane-services/llm-api-gateway/config", "//src/invocation-plane-services/llm-api-gateway/internal/ptr", "//src/invocation-plane-services/llm-api-gateway/internal/servicetier", + "//src/invocation-plane-services/llm-api-gateway/lastcluster", "//src/invocation-plane-services/llm-api-gateway/models", "//src/invocation-plane-services/llm-api-gateway/ratelimit", "//src/invocation-plane-services/llm-api-gateway/requestctx", @@ -50,6 +51,7 @@ go_test( name = "provider_test", srcs = [ "chat_passthrough_test.go", + "stargate_last_cluster_test.go", "stargate_test.go", ], embed = [":provider"], @@ -57,6 +59,7 @@ go_test( "//src/invocation-plane-services/llm-api-gateway/config", "//src/invocation-plane-services/llm-api-gateway/internal/ptr", "//src/invocation-plane-services/llm-api-gateway/internal/servicetier", + "//src/invocation-plane-services/llm-api-gateway/lastcluster", "//src/invocation-plane-services/llm-api-gateway/models", "//src/invocation-plane-services/llm-api-gateway/requestctx", "@com_github_stretchr_testify//require", diff --git a/src/invocation-plane-services/llm-api-gateway/provider/provider.go b/src/invocation-plane-services/llm-api-gateway/provider/provider.go index f6d42c8a28..ce6bd8b581 100644 --- a/src/invocation-plane-services/llm-api-gateway/provider/provider.go +++ b/src/invocation-plane-services/llm-api-gateway/provider/provider.go @@ -24,6 +24,7 @@ import ( "io" "net/http" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/lastcluster" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/models" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/ratelimit" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/requestctx" @@ -94,6 +95,12 @@ type InferenceProvider interface { ) (<-chan StreamEvent, error) } +// LastClusterAware is implemented by providers that send the session's last +// serving cluster to Stargate. +type LastClusterAware interface { + SetLastClusterTracker(tracker *lastcluster.Tracker) +} + type OpenAIProxyProvider interface { Proxy( ctx context.Context, diff --git a/src/invocation-plane-services/llm-api-gateway/provider/stargate.go b/src/invocation-plane-services/llm-api-gateway/provider/stargate.go index 3473e75070..d161ae46b7 100644 --- a/src/invocation-plane-services/llm-api-gateway/provider/stargate.go +++ b/src/invocation-plane-services/llm-api-gateway/provider/stargate.go @@ -41,6 +41,7 @@ import ( "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/config" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/internal/ptr" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/internal/servicetier" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/lastcluster" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/models" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/requestctx" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/telemetry" @@ -74,12 +75,22 @@ const ( sseMaxToken = 4 * 1024 * 1024 ) +const ( + // headerStargateClusterID names the cluster that served a request. + headerStargateClusterID = "X-Stargate-Cluster-Id" + // headerStargateLastClusterID is gateway-owned: it carries the cluster + // that served the session's last successful request. + headerStargateLastClusterID = "X-Stargate-Last-Cluster-Id" +) + type StargateProvider struct { baseURL *url.URL client *http.Client requestTimeout time.Duration upstreamRequestsTotal otelmetric.Int64Counter upstreamRequestDuration otelmetric.Float64Histogram + // lastCluster is nil when last-cluster hints are disabled. + lastCluster *lastcluster.Tracker } func NewStargateProvider(cfg config.StargateConfig) (*StargateProvider, error) { @@ -120,6 +131,36 @@ func NewStargateProvider(cfg config.StargateConfig) (*StargateProvider, error) { }, nil } +// SetLastClusterTracker enables last-cluster hints. Call it before serving +// requests. +func (p *StargateProvider) SetLastClusterTracker(tracker *lastcluster.Tracker) { + p.lastCluster = tracker +} + +// setLastClusterHeader replaces any inbound last-cluster hint with the +// gateway's own value, or removes it when there is none. +func setLastClusterHeader(headers http.Header, clusterID string) { + headers.Del(headerStargateLastClusterID) + if clusterID != "" { + headers.Set(headerStargateLastClusterID, clusterID) + } +} + +// rememberLastCluster records the serving cluster at response-header time. +// The tracker ignores non-2xx responses and writes in the background. +func (p *StargateProvider) rememberLastCluster( + ctx context.Context, + reqCtx *requestctx.RequestContext, + model string, + key string, + resp *http.Response, +) { + if key == "" || resp == nil { + return + } + p.lastCluster.Remember(ctx, reqCtx, model, key, resp.StatusCode, resp.Header.Get(headerStargateClusterID)) +} + func (p *StargateProvider) Complete( ctx context.Context, reqCtx *requestctx.RequestContext, @@ -129,6 +170,9 @@ func (p *StargateProvider) Complete( if err != nil { return nil, err } + lastClusterModel := effectiveModel(reqCtx, request) + lastClusterKey, lastClusterID := p.lastCluster.Lookup(ctx, reqCtx, lastClusterModel) + setLastClusterHeader(outbound.Header, lastClusterID) requestCtx, cancel := p.requestContext(ctx) start := time.Now() @@ -145,6 +189,7 @@ func (p *StargateProvider) Complete( p.recordUpstreamRequest(ctx, reqCtx, start, resp.StatusCode, err) return nil, err } + p.rememberLastCluster(ctx, reqCtx, lastClusterModel, lastClusterKey, resp) events := make(chan StreamEvent, 8) go p.readStream(requestCtx, cancel, resp.Body, events) @@ -169,6 +214,9 @@ func (p *StargateProvider) Stream( if err != nil { return nil, err } + lastClusterModel := effectiveModel(reqCtx, request) + lastClusterKey, lastClusterID := p.lastCluster.Lookup(ctx, reqCtx, lastClusterModel) + setLastClusterHeader(outbound.Header, lastClusterID) requestCtx, cancel := p.requestContext(ctx) start := time.Now() @@ -187,6 +235,7 @@ func (p *StargateProvider) Stream( } p.recordUpstreamRequest(ctx, reqCtx, start, resp.StatusCode, nil) + p.rememberLastCluster(ctx, reqCtx, lastClusterModel, lastClusterKey, resp) events := make(chan StreamEvent, 8) go p.readStream(requestCtx, cancel, resp.Body, events) @@ -336,6 +385,12 @@ func (p *StargateProvider) Proxy( return nil, errors.New("proxy path is required") } + var lastClusterModel, lastClusterKey, lastClusterID string + if reqCtx != nil { + lastClusterModel = reqCtx.Model + lastClusterKey, lastClusterID = p.lastCluster.Lookup(ctx, reqCtx, lastClusterModel) + } + requestCtx, cancel := p.requestContext(ctx) target := p.baseURL.ResolveReference(&url.URL{ Path: request.Path, @@ -366,6 +421,9 @@ func (p *StargateProvider) Proxy( outbound.Header.Del(headerPriority) // An inbound X-Routing-Method request header must never reach the router. outbound.Header.Del(headerRoutingMethod) + // X-Stargate-Last-Cluster-Id is gateway-owned: only the gateway store may + // set it. + setLastClusterHeader(outbound.Header, lastClusterID) if reqCtx != nil { if reqCtx.RequestID != "" { @@ -417,6 +475,7 @@ func (p *StargateProvider) Proxy( } p.recordUpstreamRequest(ctx, reqCtx, start, resp.StatusCode, nil) + p.rememberLastCluster(ctx, reqCtx, lastClusterModel, lastClusterKey, resp) return &ProxyResponse{ StatusCode: clientStatusCode(resp), Header: resp.Header.Clone(), diff --git a/src/invocation-plane-services/llm-api-gateway/provider/stargate_last_cluster_test.go b/src/invocation-plane-services/llm-api-gateway/provider/stargate_last_cluster_test.go new file mode 100644 index 0000000000..c6bd42cd60 --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/provider/stargate_last_cluster_test.go @@ -0,0 +1,543 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package provider + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/config" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/internal/ptr" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/lastcluster" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/models" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/requestctx" +) + +// lastClusterStore is an in-memory lastcluster.Store that counts calls, can +// inject errors, and can block Get while ignoring the context. +type lastClusterStore struct { + mu sync.Mutex + values map[string]string + gets int + sets int + getErr error + setErr error + blockGet chan struct{} +} + +func newLastClusterStore() *lastClusterStore { + return &lastClusterStore{values: map[string]string{}} +} + +func (s *lastClusterStore) Get(_ context.Context, key string) (string, bool, error) { + s.mu.Lock() + s.gets++ + block, err := s.blockGet, s.getErr + value, ok := s.values[key] + s.mu.Unlock() + if block != nil { + <-block + } + if err != nil { + return "", false, err + } + return value, ok, nil +} + +func (s *lastClusterStore) Set(_ context.Context, key, value string, _ time.Duration) error { + s.mu.Lock() + defer s.mu.Unlock() + s.sets++ + if s.setErr != nil { + return s.setErr + } + s.values[key] = value + return nil +} + +func (s *lastClusterStore) get(key string) string { + s.mu.Lock() + defer s.mu.Unlock() + return s.values[key] +} + +func (s *lastClusterStore) counts() (gets, sets int) { + s.mu.Lock() + defer s.mu.Unlock() + return s.gets, s.sets +} + +const lastClusterTestModel = "company/model" + +func lastClusterReqCtx(source string) *requestctx.RequestContext { + return &requestctx.RequestContext{ + RequestID: "req-lc", + RoutingKey: "fn-lc", + Model: lastClusterTestModel, + SessionID: "session-1", + SessionSource: source, + CacheAffinityKey: "mt:v1:session:lc-hash", + } +} + +func lastClusterKey(reqCtx *requestctx.RequestContext) string { + return lastcluster.Key(reqCtx.RoutingKey, reqCtx.Model, reqCtx.CacheAffinityKey) +} + +func lastClusterChatRequest() *NormalizedRequest { + return &NormalizedRequest{ + ChatRequest: &models.ChatCompletionRequest{ + Model: lastClusterTestModel, + Messages: &[]models.ChatMessage{ + {Role: models.ChatCompletionRoleUser, Content: models.SingleTextContent("hi")}, + }, + }, + } +} + +func lastClusterSSE(t *testing.T) string { + return sseChatBody(t, models.ChatCompletionChunk{ + ID: "chatcmpl-lc", + Object: models.ObjectChatCompletionChunk, + Model: lastClusterTestModel, + Choices: []models.ChatCompletionChunkChoice{{ + Delta: models.ChatCompletionChunkDelta{Content: ptr.To("ok")}, + FinishReason: ptr.To(models.FinishReasonStop), + }}, + }) +} + +func newLastClusterTracker(store lastcluster.Store) *lastcluster.Tracker { + return lastcluster.NewTracker(store, lastcluster.Options{TTL: time.Minute, LookupTimeout: 20 * time.Millisecond}) +} + +// lastClusterUpstream answers every request with status and the cluster ID +// header, and records the last-cluster hint each request carried. +type lastClusterUpstream struct { + mu sync.Mutex + hints []string + status int + clusterID string +} + +func (u *lastClusterUpstream) roundTrip(t *testing.T) roundTripFunc { + return func(r *http.Request) (*http.Response, error) { + u.mu.Lock() + u.hints = append(u.hints, r.Header.Get(headerStargateLastClusterID)) + status, clusterID := u.status, u.clusterID + u.mu.Unlock() + header := http.Header{headerContentType: []string{contentTypeSSE}} + if clusterID != "" { + header.Set(headerStargateClusterID, clusterID) + } + body := lastClusterSSE(t) + if status >= http.StatusBadRequest { + header.Set(headerContentType, contentTypeJSON) + body = `{"error":{"message":"upstream failed"}}` + } + return &http.Response{ + StatusCode: status, + Header: header, + Body: io.NopCloser(strings.NewReader(body)), + Request: r, + }, nil + } +} + +func (u *lastClusterUpstream) lastHint() string { + u.mu.Lock() + defer u.mu.Unlock() + if len(u.hints) == 0 { + return "" + } + return u.hints[len(u.hints)-1] +} + +// lastClusterCall runs one request through the named provider path and +// closes any response body. +func lastClusterCall( + ctx context.Context, + p *StargateProvider, + path string, + reqCtx *requestctx.RequestContext, + inbound http.Header, +) (int, error) { + switch path { + case "Complete": + _, err := p.Complete(ctx, reqCtx, lastClusterChatRequest()) + return 0, err + case "Stream": + events, err := p.Stream(ctx, reqCtx, lastClusterChatRequest()) + if err != nil { + return 0, err + } + for range events { + } + return 0, nil + default: + resp, err := p.Proxy(ctx, reqCtx, &ProxyRequest{ + Method: http.MethodPost, + Path: "/v1/responses", + Header: inbound, + Body: io.NopCloser(strings.NewReader(`{}`)), + }) + if err != nil { + return 0, err + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + return resp.StatusCode, nil + } +} + +var lastClusterPaths = []string{"Complete", "Stream", "Proxy"} + +func TestLastClusterRoundTrip(t *testing.T) { + for _, path := range lastClusterPaths { + for _, source := range []string{ + requestctx.SessionSourcePromptCacheKey, + requestctx.SessionSourceConversationID, + requestctx.SessionSourceHeader, + } { + t.Run(path+"/"+source, func(t *testing.T) { + ctx := context.Background() + store := newLastClusterStore() + tracker := newLastClusterTracker(store) + upstream := &lastClusterUpstream{status: http.StatusOK, clusterID: "cluster-b"} + p, err := NewStargateProvider(config.StargateConfig{URL: "http://stargate.example"}) + require.NoError(t, err) + p.client = &http.Client{Transport: upstream.roundTrip(t)} + p.SetLastClusterTracker(tracker) + reqCtx := lastClusterReqCtx(source) + + _, err = lastClusterCall(ctx, p, path, reqCtx, nil) + require.NoError(t, err) + require.Empty(t, upstream.lastHint(), "first request must carry no hint") + tracker.Drain(ctx) + require.Equal(t, "cluster-b", store.get(lastClusterKey(reqCtx))) + + _, err = lastClusterCall(ctx, p, path, reqCtx, nil) + require.NoError(t, err) + require.Equal(t, "cluster-b", upstream.lastHint()) + }) + } + } +} + +func TestLastClusterInboundHintStripped(t *testing.T) { + for _, path := range lastClusterPaths { + for _, enabled := range []bool{false, true} { + name := path + "/disabled" + if enabled { + name = path + "/enabled" + } + t.Run(name, func(t *testing.T) { + ctx := context.Background() + upstream := &lastClusterUpstream{status: http.StatusOK} + p, err := NewStargateProvider(config.StargateConfig{URL: "http://stargate.example"}) + require.NoError(t, err) + p.client = &http.Client{Transport: upstream.roundTrip(t)} + reqCtx := lastClusterReqCtx(requestctx.SessionSourceHeader) + want := "" + if enabled { + store := newLastClusterStore() + store.values[lastClusterKey(reqCtx)] = "cluster-b" + p.SetLastClusterTracker(newLastClusterTracker(store)) + want = "cluster-b" + } + inbound := http.Header{} + inbound.Set(headerStargateLastClusterID, "client-forged") + + _, err = lastClusterCall(ctx, p, path, reqCtx, inbound) + require.NoError(t, err) + require.Equal(t, want, upstream.lastHint()) + + // Ineligible session: the forged value is still dropped. + _, err = lastClusterCall(ctx, p, path, lastClusterReqCtx(requestctx.SessionSourcePayload), inbound) + require.NoError(t, err) + require.Empty(t, upstream.lastHint()) + }) + } + } +} + +func TestLastClusterDisabledProxyHeadersUnchanged(t *testing.T) { + var got []http.Header + p, err := NewStargateProvider(config.StargateConfig{URL: "http://stargate.example"}) + require.NoError(t, err) + p.client = &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) { + got = append(got, r.Header.Clone()) + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{headerStargateClusterID: []string{"cluster-b"}}, + Body: io.NopCloser(strings.NewReader(`{}`)), + Request: r, + }, nil + })} + reqCtx := lastClusterReqCtx(requestctx.SessionSourcePromptCacheKey) + baseline := http.Header{"X-Client": []string{"1"}} + withHint := baseline.Clone() + withHint.Set(headerStargateLastClusterID, "client-forged") + + for _, inbound := range []http.Header{baseline, withHint, baseline} { + _, err := lastClusterCall(context.Background(), p, "Proxy", reqCtx, inbound) + require.NoError(t, err) + } + require.Len(t, got, 3) + require.Equal(t, got[0], got[1], "inbound hint must only be stripped") + require.Equal(t, got[0], got[2], "a disabled gateway must not learn from responses") +} + +func TestLastClusterPayloadSessionSkipped(t *testing.T) { + for _, path := range lastClusterPaths { + t.Run(path, func(t *testing.T) { + ctx := context.Background() + store := newLastClusterStore() + tracker := newLastClusterTracker(store) + upstream := &lastClusterUpstream{status: http.StatusOK, clusterID: "cluster-b"} + p, err := NewStargateProvider(config.StargateConfig{URL: "http://stargate.example"}) + require.NoError(t, err) + p.client = &http.Client{Transport: upstream.roundTrip(t)} + p.SetLastClusterTracker(tracker) + + for i := 0; i < 2; i++ { + _, err := lastClusterCall(ctx, p, path, lastClusterReqCtx(requestctx.SessionSourcePayload), nil) + require.NoError(t, err) + require.Empty(t, upstream.lastHint()) + } + tracker.Drain(ctx) + gets, sets := store.counts() + require.Zero(t, gets) + require.Zero(t, sets) + }) + } +} + +func TestLastClusterKeyIsolation(t *testing.T) { + ctx := context.Background() + store := newLastClusterStore() + upstream := &lastClusterUpstream{status: http.StatusOK} + p, err := NewStargateProvider(config.StargateConfig{URL: "http://stargate.example"}) + require.NoError(t, err) + p.client = &http.Client{Transport: upstream.roundTrip(t)} + p.SetLastClusterTracker(newLastClusterTracker(store)) + seeded := lastClusterReqCtx(requestctx.SessionSourceHeader) + store.values[lastClusterKey(seeded)] = "cluster-b" + + otherModel := lastClusterReqCtx(requestctx.SessionSourceHeader) + otherModel.Model = "company/other-model" + otherRoutingKey := lastClusterReqCtx(requestctx.SessionSourceHeader) + otherRoutingKey.RoutingKey = "fn-other" + + for _, reqCtx := range []*requestctx.RequestContext{otherModel, otherRoutingKey} { + _, err := lastClusterCall(ctx, p, "Proxy", reqCtx, nil) + require.NoError(t, err) + require.Empty(t, upstream.lastHint()) + } + _, err = lastClusterCall(ctx, p, "Proxy", seeded, nil) + require.NoError(t, err) + require.Equal(t, "cluster-b", upstream.lastHint()) +} + +func TestLastClusterNoWriteOnFailure(t *testing.T) { + for _, path := range lastClusterPaths { + t.Run(path, func(t *testing.T) { + ctx := context.Background() + store := newLastClusterStore() + tracker := newLastClusterTracker(store) + reqCtx := lastClusterReqCtx(requestctx.SessionSourceHeader) + store.values[lastClusterKey(reqCtx)] = "cluster-a" + p, err := NewStargateProvider(config.StargateConfig{URL: "http://stargate.example"}) + require.NoError(t, err) + p.SetLastClusterTracker(tracker) + + for _, status := range []int{http.StatusServiceUnavailable, http.StatusTooManyRequests, http.StatusBadRequest} { + upstream := &lastClusterUpstream{status: status, clusterID: "cluster-c"} + p.client = &http.Client{Transport: upstream.roundTrip(t)} + _, _ = lastClusterCall(ctx, p, path, reqCtx, nil) + require.Equal(t, "cluster-a", upstream.lastHint()) + } + p.client = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("connection reset") + })} + _, err = lastClusterCall(ctx, p, path, reqCtx, nil) + require.Error(t, err) + + tracker.Drain(ctx) + require.Equal(t, "cluster-a", store.get(lastClusterKey(reqCtx))) + _, sets := store.counts() + require.Zero(t, sets) + }) + } +} + +func TestLastClusterStreamWritesBeforeEnd(t *testing.T) { + for _, path := range []string{"Stream", "Proxy"} { + t.Run(path, func(t *testing.T) { + release := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set(headerContentType, contentTypeSSE) + w.Header().Set(headerStargateClusterID, "cluster-b") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, "data: {\"id\":\"c\",\"object\":\"chat.completion.chunk\",\"choices\":[]}\n\n") + w.(http.Flusher).Flush() + select { + case <-release: + case <-r.Context().Done(): + } + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer server.Close() + var releaseOnce sync.Once + releaseStream := func() { releaseOnce.Do(func() { close(release) }) } + defer releaseStream() + + ctx := context.Background() + store := newLastClusterStore() + tracker := newLastClusterTracker(store) + p, err := NewStargateProvider(config.StargateConfig{URL: server.URL}) + require.NoError(t, err) + p.SetLastClusterTracker(tracker) + reqCtx := lastClusterReqCtx(requestctx.SessionSourcePromptCacheKey) + + var body io.Closer + if path == "Stream" { + events, err := p.Stream(ctx, reqCtx, lastClusterChatRequest()) + require.NoError(t, err) + <-events + defer func() { + releaseStream() + for range events { + } + }() + } else { + resp, err := p.Proxy(ctx, reqCtx, &ProxyRequest{ + Method: http.MethodPost, + Path: "/v1/responses", + Body: io.NopCloser(strings.NewReader(`{}`)), + }) + require.NoError(t, err) + body = resp.Body + defer body.Close() + } + + require.Eventually(t, func() bool { + return store.get(lastClusterKey(reqCtx)) == "cluster-b" + }, 2*time.Second, 5*time.Millisecond, "write must land while the stream is still open") + }) + } +} + +func TestLastClusterLookupTimeoutProceeds(t *testing.T) { + for _, path := range lastClusterPaths { + t.Run(path, func(t *testing.T) { + store := newLastClusterStore() + store.blockGet = make(chan struct{}) + defer close(store.blockGet) + p, err := NewStargateProvider(config.StargateConfig{URL: "http://stargate.example"}) + require.NoError(t, err) + p.SetLastClusterTracker(newLastClusterTracker(store)) + var arrived time.Time + upstream := &lastClusterUpstream{status: http.StatusOK} + inner := upstream.roundTrip(t) + p.client = &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) { + arrived = time.Now() + return inner(r) + })} + + start := time.Now() + _, err = lastClusterCall(context.Background(), p, path, lastClusterReqCtx(requestctx.SessionSourceHeader), nil) + require.NoError(t, err) + require.Less(t, arrived.Sub(start), 25*time.Millisecond) + require.Empty(t, upstream.lastHint()) + }) + } +} + +func TestLastClusterStoreErrorsTransparent(t *testing.T) { + for _, path := range lastClusterPaths { + t.Run(path, func(t *testing.T) { + ctx := context.Background() + // run returns everything the client would see: the error, plus the + // proxy status and body, the aggregated completion, or every stream + // event. + run := func(tracker *lastcluster.Tracker) (status int, err error, output string) { + p, newErr := NewStargateProvider(config.StargateConfig{URL: "http://stargate.example"}) + require.NoError(t, newErr) + p.SetLastClusterTracker(tracker) + upstream := &lastClusterUpstream{status: http.StatusOK, clusterID: "cluster-b"} + p.client = &http.Client{Transport: upstream.roundTrip(t)} + reqCtx := lastClusterReqCtx(requestctx.SessionSourceHeader) + defer tracker.Drain(ctx) + + switch path { + case "Complete": + response, err := p.Complete(ctx, reqCtx, lastClusterChatRequest()) + raw, marshalErr := json.Marshal(response) + require.NoError(t, marshalErr) + return 0, err, string(raw) + case "Stream": + events, err := p.Stream(ctx, reqCtx, lastClusterChatRequest()) + if err != nil { + return 0, err, "" + } + var seen []string + for event := range events { + raw, marshalErr := json.Marshal(event.Chunk) + require.NoError(t, marshalErr) + seen = append(seen, fmt.Sprintf("%s err=%v", raw, event.Err)) + } + return 0, nil, strings.Join(seen, "\n") + default: + resp, err := p.Proxy(ctx, reqCtx, &ProxyRequest{ + Method: http.MethodPost, Path: "/v1/responses", Body: io.NopCloser(strings.NewReader(`{}`)), + }) + require.NoError(t, err) + raw, _ := io.ReadAll(resp.Body) + _ = resp.Body.Close() + return resp.StatusCode, nil, string(raw) + } + } + + store := newLastClusterStore() + store.getErr = errors.New("olric unavailable") + store.setErr = errors.New("olric unavailable") + gotStatus, gotErr, gotOutput := run(newLastClusterTracker(store)) + wantStatus, wantErr, wantOutput := run(nil) + + require.Equal(t, wantStatus, gotStatus) + require.Equal(t, wantErr, gotErr) + require.NotEmpty(t, wantOutput) + require.Equal(t, wantOutput, gotOutput) + gets, sets := store.counts() + require.Equal(t, 1, gets) + require.Equal(t, 1, sets) + }) + } +} diff --git a/src/invocation-plane-services/llm-api-gateway/requestctx/requestctx.go b/src/invocation-plane-services/llm-api-gateway/requestctx/requestctx.go index 8e6d2ec17f..4c64426249 100644 --- a/src/invocation-plane-services/llm-api-gateway/requestctx/requestctx.go +++ b/src/invocation-plane-services/llm-api-gateway/requestctx/requestctx.go @@ -21,6 +21,15 @@ import ( "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/nvcf" ) +// Session sources recorded in RequestContext.SessionSource. +const ( + SessionSourceHeader = "header" + SessionSourcePayload = "payload" + SessionSourcePromptCacheKey = "prompt_cache_key" + SessionSourceConversationID = "conversation_id" + SessionSourceClaudeCodeHeader = "claude_code_header" +) + type RequestContext struct { RequestID string APIKeyID string // client auth subject; reserved for attribution, auditing, or future API-key policy hooks diff --git a/src/invocation-plane-services/llm-api-gateway/server/BUILD.bazel b/src/invocation-plane-services/llm-api-gateway/server/BUILD.bazel index b7acba0879..1d1ec795d2 100644 --- a/src/invocation-plane-services/llm-api-gateway/server/BUILD.bazel +++ b/src/invocation-plane-services/llm-api-gateway/server/BUILD.bazel @@ -23,6 +23,7 @@ go_library( deps = [ "//src/invocation-plane-services/llm-api-gateway/api", "//src/invocation-plane-services/llm-api-gateway/config", + "//src/invocation-plane-services/llm-api-gateway/lastcluster", "//src/invocation-plane-services/llm-api-gateway/provider", "//src/invocation-plane-services/llm-api-gateway/ratelimit", "//src/invocation-plane-services/llm-api-gateway/ratelimitsync", @@ -41,10 +42,14 @@ alias( go_test( name = "server_test", - srcs = ["server_test.go"], + srcs = [ + "last_cluster_test.go", + "server_test.go", + ], embed = [":server"], deps = [ "//src/invocation-plane-services/llm-api-gateway/config", + "//src/invocation-plane-services/llm-api-gateway/internal/olrictest", "//src/invocation-plane-services/llm-api-gateway/provider", ], ) diff --git a/src/invocation-plane-services/llm-api-gateway/server/last_cluster_test.go b/src/invocation-plane-services/llm-api-gateway/server/last_cluster_test.go new file mode 100644 index 0000000000..6f85af7ee5 --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/server/last_cluster_test.go @@ -0,0 +1,177 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package server + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/internal/olrictest" +) + +// hintingStargate answers like Stargate with x-stargate-cluster-id set and +// records the x-stargate-last-cluster-id each request carried. +type hintingStargate struct { + *httptest.Server + mu sync.Mutex + hints []string +} + +func newHintingStargate(t *testing.T, clusterID string) *hintingStargate { + t.Helper() + s := &hintingStargate{} + s.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + s.mu.Lock() + s.hints = append(s.hints, r.Header.Get("X-Stargate-Last-Cluster-Id")) + s.mu.Unlock() + w.Header().Set("X-Stargate-Cluster-Id", clusterID) + switch r.URL.Path { + case "/v1/chat/completions": + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"company-name/model-name","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]}`+"\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + case "/v1/embeddings": + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1]}],"model":"company-name/model-name","usage":{"prompt_tokens":1,"total_tokens":1}}`) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(s.Close) + return s +} + +func (s *hintingStargate) lastHint() string { + s.mu.Lock() + defer s.mu.Unlock() + if len(s.hints) == 0 { + return "" + } + return s.hints[len(s.hints)-1] +} + +const lastClusterChatPayload = `{"model":"fn-alpha/company-name/model-name","prompt_cache_key":"session-1","messages":[{"role":"user","content":"hello"}]}` + +// assertSecondRequestCarriesHint sends a first request, then repeats the +// session's request until the background write is visible. +func assertSecondRequestCarriesHint(t *testing.T, gateway *testGateway, upstream *hintingStargate) { + t.Helper() + status, body, err := gateway.post("/v1/chat/completions", lastClusterChatPayload) + if err != nil || status != http.StatusOK { + t.Fatalf("first request status=%d err=%v body=%s", status, err, body) + } + if hint := upstream.lastHint(); hint != "" { + t.Fatalf("first request hint = %q, want none", hint) + } + deadline := time.Now().Add(5 * time.Second) + for { + status, body, err := gateway.post("/v1/chat/completions", lastClusterChatPayload) + if err != nil || status != http.StatusOK { + t.Fatalf("follow-up request status=%d err=%v body=%s", status, err, body) + } + if upstream.lastHint() == "cluster-b" { + return + } + if time.Now().After(deadline) { + t.Fatalf("follow-up hint = %q, want cluster-b", upstream.lastHint()) + } + time.Sleep(10 * time.Millisecond) + } +} + +func TestOlricWithoutRateLimiterServesLastCluster(t *testing.T) { + if testing.Short() { + t.Skip("skipping embedded Olric test in -short mode") + } + t.Setenv("POD_NAMESPACE", "") + upstream := newHintingStargate(t, "cluster-b") + cfg := testConfig(upstream.URL) + cfg.RateLimiter.Enabled = false + cfg.Olric = olrictest.NodeConfig(t, nil) + cfg.Stargate.LastCluster.Enabled = true + + assertSecondRequestCarriesHint(t, startGatewayWithConfig(t, cfg, protocols[0]), upstream) +} + +func TestLocalStoreServesLastClusterWhenOlricOff(t *testing.T) { + upstream := newHintingStargate(t, "cluster-b") + cfg := testConfig(upstream.URL) + cfg.Olric.Enabled = false + cfg.Stargate.LastCluster.Enabled = true + + assertSecondRequestCarriesHint(t, startGatewayWithConfig(t, cfg, protocols[0]), upstream) +} + +func TestLastClusterDisabledSendsNoHint(t *testing.T) { + upstream := newHintingStargate(t, "cluster-b") + cfg := testConfig(upstream.URL) + gateway := startGatewayWithConfig(t, cfg, protocols[0]) + + for i := 0; i < 3; i++ { + status, body, err := gateway.post("/v1/chat/completions", lastClusterChatPayload) + if err != nil || status != http.StatusOK { + t.Fatalf("request status=%d err=%v body=%s", status, err, body) + } + if hint := upstream.lastHint(); hint != "" { + t.Fatalf("request %d hint = %q, want none while disabled", i, hint) + } + } +} + +func TestStargateClusterIDStrippedFromClientResponses(t *testing.T) { + upstream := newHintingStargate(t, "cluster-b") + cfg := testConfig(upstream.URL) + cfg.Stargate.LastCluster.Enabled = true + gateway := startGatewayWithConfig(t, cfg, protocols[0]) + + for _, tc := range []struct{ path, payload string }{ + {"/v1/chat/completions", lastClusterChatPayload}, + {"/v1/chat/completions", `{"model":"fn-alpha/company-name/model-name","prompt_cache_key":"s","stream":true,"messages":[{"role":"user","content":"hello"}]}`}, + {"/v1/embeddings", `{"model":"fn-alpha/company-name/model-name","input":"hello"}`}, + } { + req, err := http.NewRequest(http.MethodPost, gateway.url+tc.path, strings.NewReader(tc.payload)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Stargate-Last-Cluster-Id", "client-forged") + resp, err := gateway.client.Do(req) + if err != nil { + t.Fatal(err) + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("%s status = %d, want 200", tc.path, resp.StatusCode) + } + if got := resp.Header.Get("X-Stargate-Cluster-Id"); got != "" { + t.Fatalf("%s leaked x-stargate-cluster-id %q to the client", tc.path, got) + } + if hint := upstream.lastHint(); hint == "client-forged" { + t.Fatalf("%s forwarded the client-supplied last-cluster hint", tc.path) + } + } +} diff --git a/src/invocation-plane-services/llm-api-gateway/server/server.go b/src/invocation-plane-services/llm-api-gateway/server/server.go index 9f37226860..658883a479 100644 --- a/src/invocation-plane-services/llm-api-gateway/server/server.go +++ b/src/invocation-plane-services/llm-api-gateway/server/server.go @@ -28,6 +28,7 @@ import ( "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/api" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/config" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/lastcluster" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/provider" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/ratelimit" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/ratelimitsync" @@ -65,10 +66,47 @@ func New( }) } - limiter, err := newRateLimiter(cfg, e) + // context.Background for startup: Echo gives us no hook-scoped ctx, and a + // startup failure cannot be undone by callers. The telemetry logger is + // initialised per-call, so ctx only affects tracing and logs. + ctx := context.Background() + olric, err := startOlric(ctx, cfg) + if err != nil { + return nil, err + } + + limiter, stopLimiter, err := newRateLimiter(cfg, olric.node) + if err != nil { + olric.stop(ctx) + return nil, err + } + + tracker, err := newLastClusterTracker(ctx, cfg, olric.node) if err != nil { + stopLimiter() + olric.stop(ctx) return nil, err } + if aware, ok := inferenceProvider.(provider.LastClusterAware); ok && tracker != nil { + aware.SetLastClusterTracker(tracker) + } + + if olric.node != nil || tracker != nil { + // One hook keeps the order fixed: stop the limiter, wait briefly for + // pending last-cluster writes, then leave the Olric cluster. This is + // best effort: http.Server.Shutdown does not wait for its hooks, so + // writes still pending when the process exits are dropped. A lost + // hint only costs that session one affinity miss. + e.Server.RegisterOnShutdown(func() { + stopLimiter() + if tracker != nil { + drainCtx, cancel := contextWithOptionalTimeout(cfg.Olric.ShutdownTimeout) + tracker.Drain(drainCtx) + cancel() + } + olric.stop(context.Background()) + }) + } handlers := api.NewHandlers( cfg, @@ -88,38 +126,61 @@ func Start(e *echo.Echo, addr string) error { return e.Server.ListenAndServe() } -func newRateLimiter(cfg *config.Config, e *echo.Echo) (ratelimit.RateLimiter, error) { - if !cfg.RateLimiter.Enabled { - return ratelimit.AllowAll, nil - } +// olricRuntime is the embedded Olric node shared by the rate limiter and the +// last-cluster store. node is nil when Olric is disabled. +type olricRuntime struct { + node *util.OlricNode + collector *telemetry.OlricCollector + timeout time.Duration +} +// startOlric starts the embedded Olric node whenever Olric is enabled, even +// when the rate limiter is off, so other gateway state can use it. +func startOlric(ctx context.Context, cfg *config.Config) (*olricRuntime, error) { + runtime := &olricRuntime{timeout: cfg.Olric.ShutdownTimeout} if !cfg.Olric.Enabled { - if cfg.RateLimiter.FailOpen { - return ratelimit.AllowAll, nil - } - return ratelimit.RejectAll, nil + return runtime, nil } - - // context.Background for startup: Echo gives us no hook-scoped ctx, and a - // rate-limiter startup failure cannot be undone by callers. The telemetry - // logger is initialised per-call, so ctx only affects tracing and logs. - ctx := context.Background() node, err := util.NewOlricNode(ctx, cfg.Olric) if err != nil { return nil, fmt.Errorf("start olric node: %w", err) } - - olricCollector, err := telemetry.NewOlricCollector(node.Client, node.SelfAddr) + collector, err := telemetry.NewOlricCollector(node.Client, node.SelfAddr) if err != nil { util.ShutdownOlricNode(ctx, node, cfg.Olric.ShutdownTimeout) return nil, fmt.Errorf("start olric metrics collector: %w", err) } + runtime.node = node + runtime.collector = collector + return runtime, nil +} + +func (r *olricRuntime) stop(ctx context.Context) { + if r == nil || r.node == nil { + return + } + r.collector.Stop() + util.ShutdownOlricNode(ctx, r.node, r.timeout) +} + +// newRateLimiter builds the limiter on the shared Olric node. The returned +// stop function is never nil. +func newRateLimiter(cfg *config.Config, node *util.OlricNode) (ratelimit.RateLimiter, func(), error) { + noop := func() {} + if !cfg.RateLimiter.Enabled { + return ratelimit.AllowAll, noop, nil + } + + if node == nil { + if cfg.RateLimiter.FailOpen { + return ratelimit.AllowAll, noop, nil + } + return ratelimit.RejectAll, noop, nil + } syncRuntime, err := ratelimitsync.NewPublisherRuntime(cfg) if err != nil { - olricCollector.Stop() - util.ShutdownOlricNode(ctx, node, cfg.Olric.ShutdownTimeout) - return nil, err + return nil, nil, err } limiter, err := ratelimit.NewRateLimiter( @@ -128,22 +189,16 @@ func newRateLimiter(cfg *config.Config, e *echo.Echo) (ratelimit.RateLimiter, er ratelimit.WithSynchronizer(syncRuntime.Synchronizer), ) if err != nil { - olricCollector.Stop() syncRuntime.Stop() - util.ShutdownOlricNode(ctx, node, cfg.Olric.ShutdownTimeout) - return nil, err + return nil, nil, err } if err := syncRuntime.Start(); err != nil { - olricCollector.Stop() syncRuntime.Stop() - util.ShutdownOlricNode(ctx, node, cfg.Olric.ShutdownTimeout) - return nil, err + return nil, nil, err } - e.Server.RegisterOnShutdown(func() { - olricCollector.Stop() - + stop := func() { // The sync synchronizer's Stop() blocks on the publisher loop draining // its in-flight publish goroutines; bound it so a stuck remote cannot // delay the gateway's shutdown indefinitely. We reuse the Olric @@ -161,10 +216,57 @@ func newRateLimiter(cfg *config.Config, e *echo.Echo) (ratelimit.RateLimiter, er Dur("timeout", cfg.Olric.ShutdownTimeout). Msg("rate limit sync runtime did not stop within shutdown timeout") } - util.ShutdownOlricNode(context.Background(), node, cfg.Olric.ShutdownTimeout) - }) + } + return limiter, stop, nil +} + +// newLastClusterTracker returns nil when last-cluster hints are disabled. It +// uses the shared Olric DMap when Olric is enabled, and an in-process LRU +// otherwise. +func newLastClusterTracker( + ctx context.Context, + cfg *config.Config, + node *util.OlricNode, +) (*lastcluster.Tracker, error) { + lcCfg := cfg.Stargate.LastCluster + if !lcCfg.Enabled { + return nil, nil + } + + var ( + store lastcluster.Store + storeKind string + ) + if node != nil { + dm, err := node.Client.NewDMap(lastcluster.DMapName) + if err != nil { + return nil, fmt.Errorf("create olric dmap %q: %w", lastcluster.DMapName, err) + } + store, storeKind = lastcluster.NewOlricStore(dm), "olric" + } else { + store, storeKind = lastcluster.NewLocalStore(lcCfg.LocalMaxEntries), "local" + } - return limiter, nil + telemetry.Logger(ctx).Info(). + Str("store", storeKind). + Dur("ttl", lcCfg.TTL). + Dur("lookup_timeout", lcCfg.LookupTimeout). + Int("local_max_entries", lcCfg.LocalMaxEntries). + Msg("stargate last-cluster hints enabled") + + return lastcluster.NewTracker(store, lastcluster.Options{ + TTL: lcCfg.TTL, + LookupTimeout: lcCfg.LookupTimeout, + }), nil +} + +// contextWithOptionalTimeout returns a context without a deadline when +// timeout <= 0. +func contextWithOptionalTimeout(timeout time.Duration) (context.Context, context.CancelFunc) { + if timeout <= 0 { + return context.WithCancel(context.Background()) + } + return context.WithTimeout(context.Background(), timeout) } // timeAfterOrForever returns a channel that never fires when timeout <= 0, diff --git a/src/invocation-plane-services/llm-api-gateway/telemetry/metrics.go b/src/invocation-plane-services/llm-api-gateway/telemetry/metrics.go index 9979f9f388..d922dc3849 100644 --- a/src/invocation-plane-services/llm-api-gateway/telemetry/metrics.go +++ b/src/invocation-plane-services/llm-api-gateway/telemetry/metrics.go @@ -98,6 +98,8 @@ func InitializeMetrics() { _ = authRequestsTotal() _ = authRequestDuration() preInitAuthMetrics() + _ = lastClusterLookupDuration() + preInitLastClusterMetrics() } func Add( @@ -466,6 +468,85 @@ func preInitAuthMetrics() { } } +// Result labels for the Stargate last-cluster store metrics. +const ( + LastClusterLookupHit = "hit" + LastClusterLookupMiss = "miss" + LastClusterLookupError = "error" + LastClusterLookupTimeout = "timeout" + LastClusterLookupSkipped = "skipped" + + LastClusterWriteOK = "ok" + LastClusterWriteError = "error" +) + +var ( + lastClusterLookupResults = []string{ + LastClusterLookupHit, + LastClusterLookupMiss, + LastClusterLookupError, + LastClusterLookupTimeout, + LastClusterLookupSkipped, + } + lastClusterWriteResults = []string{LastClusterWriteOK, LastClusterWriteError} +) + +// LastClusterLookupBuckets cover the store lookup latency. The lookup runs on +// the request path with a 20ms default deadline, so buckets stop at 50ms. +var LastClusterLookupBuckets = []float64{0.001, 0.002, 0.005, 0.01, 0.02, 0.05} + +func lastClusterLookupsTotal() otelmetric.Int64Counter { + return must.Get(Meter().Int64Counter( + metricPrefix+"last_cluster_lookups_total", + otelmetric.WithDescription("Stargate last-cluster store lookups, labelled by result."), + )) +} + +func lastClusterWritesTotal() otelmetric.Int64Counter { + return must.Get(Meter().Int64Counter( + metricPrefix+"last_cluster_writes_total", + otelmetric.WithDescription("Stargate last-cluster store writes, labelled by result."), + )) +} + +func lastClusterLookupDuration() otelmetric.Float64Histogram { + return must.Get(Meter().Float64Histogram( + metricPrefix+"last_cluster_lookup_duration_seconds", + otelmetric.WithUnit("s"), + otelmetric.WithDescription("Duration of Stargate last-cluster store lookups on the request path."), + otelmetric.WithExplicitBucketBoundaries(LastClusterLookupBuckets...), + )) +} + +// RecordLastClusterLookup counts one lookup. Skipped lookups never touch the +// store, so they are counted but not timed. +func RecordLastClusterLookup(ctx context.Context, result string, elapsed time.Duration) { + attrs := otelmetric.WithAttributes(attribute.String("result", result)) + lastClusterLookupsTotal().Add(ctx, 1, attrs) + if result != LastClusterLookupSkipped { + lastClusterLookupDuration().Record(ctx, elapsed.Seconds(), attrs) + } +} + +func RecordLastClusterWrite(ctx context.Context, result string) { + lastClusterWritesTotal().Add(ctx, 1, otelmetric.WithAttributes(attribute.String("result", result))) +} + +// preInitLastClusterMetrics emits a zero sample for every result label so +// absent() and rate() work on the first scrape, whether or not the feature +// is enabled. +func preInitLastClusterMetrics() { + ctx := context.Background() + lookups := lastClusterLookupsTotal() + for _, result := range lastClusterLookupResults { + lookups.Add(ctx, 0, otelmetric.WithAttributes(attribute.String("result", result))) + } + writes := lastClusterWritesTotal() + for _, result := range lastClusterWriteResults { + writes.Add(ctx, 0, otelmetric.WithAttributes(attribute.String("result", result))) + } +} + func DropReasonSameCluster() attribute.KeyValue { return attribute.String("reason", dropReasonSameCluster) } diff --git a/src/invocation-plane-services/llm-api-gateway/telemetry/metrics_test.go b/src/invocation-plane-services/llm-api-gateway/telemetry/metrics_test.go index 5bf3e0c795..e282b37d13 100644 --- a/src/invocation-plane-services/llm-api-gateway/telemetry/metrics_test.go +++ b/src/invocation-plane-services/llm-api-gateway/telemetry/metrics_test.go @@ -104,6 +104,8 @@ func TestMetricsDefinitionsUseServiceScopedNames(t *testing.T) { "llm_api_gateway_rate_limit_synchronizer_queue_wait_seconds", "llm_api_gateway_rate_limit_synchronizer_queue_length", "llm_api_gateway_rate_limit_synchronizer_events_dropped_total", + "llm_api_gateway_last_cluster_lookups_total", + "llm_api_gateway_last_cluster_writes_total", } { if !names[want] { t.Fatalf("missing metric %q in %#v", want, names) @@ -111,6 +113,99 @@ func TestMetricsDefinitionsUseServiceScopedNames(t *testing.T) { } } +func TestLastClusterMetricsPreInitialized(t *testing.T) { + reader := installManualReader(t) + + InitializeMetrics() + + var rm metricdata.ResourceMetrics + if err := reader.Collect(context.Background(), &rm); err != nil { + t.Fatalf("Collect() error = %v", err) + } + want := map[string][]string{ + "llm_api_gateway_last_cluster_lookups_total": {"hit", "miss", "error", "timeout", "skipped"}, + "llm_api_gateway_last_cluster_writes_total": {"ok", "error"}, + } + for name, results := range want { + sum, ok := findInt64Sum(rm, name) + if !ok { + t.Fatalf("metric %q not found", name) + } + got := map[string]int64{} + for _, dp := range sum.DataPoints { + value, _ := dp.Attributes.Value("result") + got[value.AsString()] = dp.Value + } + if len(got) != len(results) { + t.Fatalf("%s results = %v, want %v", name, got, results) + } + for _, result := range results { + value, ok := got[result] + if !ok || value != 0 { + t.Fatalf("%s{result=%q} = %d (present=%v), want 0", name, result, value, ok) + } + } + } +} + +func TestLastClusterLookupDurationBuckets(t *testing.T) { + reader := installManualReader(t) + + RecordLastClusterLookup(context.Background(), LastClusterLookupHit, 3*time.Millisecond) + RecordLastClusterLookup(context.Background(), LastClusterLookupSkipped, 0) + + var rm metricdata.ResourceMetrics + if err := reader.Collect(context.Background(), &rm); err != nil { + t.Fatalf("Collect() error = %v", err) + } + for _, sm := range rm.ScopeMetrics { + for _, m := range sm.Metrics { + if m.Name != "llm_api_gateway_last_cluster_lookup_duration_seconds" { + continue + } + hist, ok := m.Data.(metricdata.Histogram[float64]) + if !ok { + t.Fatalf("lookup duration data = %T, want float64 histogram", m.Data) + } + if len(hist.DataPoints) != 1 { + t.Fatalf("lookup duration points = %d, want 1 (skipped is not timed)", len(hist.DataPoints)) + } + bounds := hist.DataPoints[0].Bounds + if fmt.Sprint(bounds) != fmt.Sprint([]float64{0.001, 0.002, 0.005, 0.01, 0.02, 0.05}) { + t.Fatalf("lookup duration bounds = %v, want 1ms to 50ms", bounds) + } + return + } + } + t.Fatal("lookup duration histogram not found") +} + +func installManualReader(t *testing.T) *sdkmetric.ManualReader { + t.Helper() + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + oldProvider := otel.GetMeterProvider() + otel.SetMeterProvider(provider) + t.Cleanup(func() { + otel.SetMeterProvider(oldProvider) + _ = provider.Shutdown(context.Background()) + }) + return reader +} + +func findInt64Sum(rm metricdata.ResourceMetrics, name string) (metricdata.Sum[int64], bool) { + for _, sm := range rm.ScopeMetrics { + for _, m := range sm.Metrics { + if m.Name != name { + continue + } + sum, ok := m.Data.(metricdata.Sum[int64]) + return sum, ok + } + } + return metricdata.Sum[int64]{}, false +} + func TestFunctionIDAttribute(t *testing.T) { tests := []struct { name string diff --git a/src/libraries/rust/stargate/docs/api-gateway-contract.md b/src/libraries/rust/stargate/docs/api-gateway-contract.md index 96060bba57..73cd7f3387 100644 --- a/src/libraries/rust/stargate/docs/api-gateway-contract.md +++ b/src/libraries/rust/stargate/docs/api-gateway-contract.md @@ -140,10 +140,27 @@ Optional trusted headers: | `x-request-slo-ms` | Per-request LB latency hint. | | `x-max-wait-ms` | Wait budget for temporarily infeasible candidates. | | `x-stargate-max-wait-ms` | Stargate internal retry budget. | +| `x-stargate-last-cluster-id` | The `x-stargate-cluster-id` from the session's last 2xx response. Used only by algorithms with `last_cluster_affinity`. Stargate consumes it and never forwards it to pylon. | The gateway must synthesize or validate these headers. Do not pass public caller-supplied routing headers through blindly. +Last-cluster rules for the gateway: + +- A session is one `x-routing-key`, `x-model`, and `x-cache-affinity-key` + tuple. +- Always delete an inbound `x-stargate-last-cluster-id` before building the + Stargate request, including on paths that clone inbound headers. +- Set the header only from the gateway's own store, and only for sessions with + a stable key. Skip sessions whose affinity key is derived from the payload, + because it changes every turn. +- Store the `x-stargate-cluster-id` of a 2xx response when the response + headers arrive, not at stream end. Do not store on non-2xx responses or + transport errors. +- Keep the lookup bounded and off the critical path: on a miss, store error, + or timeout, send the request without the header. +- Writes must never delay or fail the client response. + Internal headers: - `x-stargate-expected-queue-ms`: Stargate-to-pylon only. Stargate strips @@ -241,3 +258,5 @@ Gateway rules: 6. Enforce streaming body rules for chat and Responses. 7. Treat `ListModels` as a short-lived hint. 8. Keep backend response headers internal unless intentionally exposed. +9. Strip inbound `x-stargate-last-cluster-id`; set it only from the gateway + store. From 85661f8db41fe8043076440b0bd41b8c14913d2c Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 14:58:55 -0700 Subject: [PATCH 3/6] test(stargate): add last-cluster end-to-end check with the LLM API gateway Add scripts/e2e-last-cluster.sh, which runs the LLM API gateway against a real Stargate, two Pylons, and two mock-dynamo backends on loopback and reproduces the overflow cache-thrash scenario for last_cluster_affinity. With the feature on, the session's 10 follow-up requests stay on the overflow cluster; with it off, they return to rank 1. It covers wait-and-widen and pulsar-wait-and-widen and exits non-zero on any mismatch. Refs #2289 Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- src/libraries/rust/stargate/scripts/README.md | 26 + .../rust/stargate/scripts/e2e-last-cluster.sh | 460 ++++++++++++++++++ 2 files changed, 486 insertions(+) create mode 100755 src/libraries/rust/stargate/scripts/e2e-last-cluster.sh diff --git a/src/libraries/rust/stargate/scripts/README.md b/src/libraries/rust/stargate/scripts/README.md index 3fd0e4cd8c..0cc03662e3 100644 --- a/src/libraries/rust/stargate/scripts/README.md +++ b/src/libraries/rust/stargate/scripts/README.md @@ -55,3 +55,29 @@ Run the support-script tests after changing repository tooling: ```bash python3 -m unittest discover -s scripts -p 'test_*.py' ``` + +## Last-Cluster End-to-End Check + +`scripts/e2e-last-cluster.sh` runs the LLM API gateway against a real +Stargate, two Pylons, and two mock-dynamo backends on loopback. It reproduces +the overflow cache-thrash scenario for `last_cluster_affinity`: + +1. Two long requests fill the session's rank-1 cluster (A). +2. The session's first gateway request overflows to cluster B. +3. A drains. +4. The next 10 gateway requests with the same `prompt_cache_key` stay on B + with the feature on, and return to A with it off. + +```bash +scripts/e2e-last-cluster.sh +scripts/e2e-last-cluster.sh --algorithm pulsar-wait-and-widen --mode on +``` + +Mode `on` sets `STARGATE_LAST_CLUSTER_ENABLED=true` in the gateway and +`last_cluster_affinity: true` in Stargate. Mode `off` disables both. The +script builds release binaries and the gateway first unless `--no-build` is +set. It attributes each request to a cluster through the mock-dynamo +`/test-control` counters, checks +`stargate_routing_session_selections_total`, and exits non-zero on any +mismatch. It needs `cargo`, `go`, `curl`, and `jq`, uses 32 loopback ports +starting at `E2E_PORT_BASE` (default `27100`), and keeps logs on failure. diff --git a/src/libraries/rust/stargate/scripts/e2e-last-cluster.sh b/src/libraries/rust/stargate/scripts/e2e-last-cluster.sh new file mode 100755 index 0000000000..fafec82e03 --- /dev/null +++ b/src/libraries/rust/stargate/scripts/e2e-last-cluster.sh @@ -0,0 +1,460 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -euo pipefail + +usage() { + cat <<'EOF' +Usage: scripts/e2e-last-cluster.sh [--algorithm NAME]... [--mode on|off]... [--no-build] + +Runs the LLM API gateway against a real Stargate, two Pylons, and two +mock-dynamo backends on loopback, and reproduces the overflow cache-thrash +scenario for last_cluster_affinity: + + 1. Session S ranks cluster A first and cluster B second. + 2. Two long requests fill A (Pylon max engine concurrency 2). + 3. The session's first gateway request overflows to B. + 4. A drains. + 5. The next 10 gateway requests in the session (same prompt_cache_key) + must stay on B with the flag on, and return to A with the flag off. + +Mode "on" enables STARGATE_LAST_CLUSTER_ENABLED in the gateway and +last_cluster_affinity in Stargate. Mode "off" disables both. Each +algorithm and mode pair runs against a fresh set of processes. The script +exits non-zero when any expectation fails. + +Options: + --algorithm NAME wait-and-widen or pulsar-wait-and-widen. Repeatable. + Default: both. + --mode MODE on or off. Repeatable. Default: both. + --no-build Reuse existing release binaries and gateway build. + +Environment: + E2E_PORT_BASE First loopback port, default 27100. Uses 32 ports. + E2E_WORK_DIR Logs and configs, default a new mktemp directory. + E2E_KEEP_WORK_DIR Set to 1 to keep the work directory after success. + +Requires cargo, go, curl, and jq. +EOF +} + +stargate_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +gateway_root="$(cd "${stargate_root}/../../../invocation-plane-services/llm-api-gateway" && pwd)" +release_dir="${stargate_root}/target/release" + +algorithms=() +modes=() +build=1 +while [[ $# -gt 0 ]]; do + case "$1" in + --algorithm) algorithms+=("$2"); shift 2 ;; + --mode) modes+=("$2"); shift 2 ;; + --no-build) build=0; shift ;; + -h|--help) usage; exit 0 ;; + *) echo "unknown argument: $1" >&2; usage >&2; exit 2 ;; + esac +done +[[ ${#algorithms[@]} -gt 0 ]] || algorithms=(wait-and-widen pulsar-wait-and-widen) +[[ ${#modes[@]} -gt 0 ]] || modes=(on off) +for algorithm in "${algorithms[@]}"; do + case "${algorithm}" in + wait-and-widen|pulsar-wait-and-widen) ;; + *) echo "unsupported algorithm: ${algorithm}" >&2; exit 2 ;; + esac +done +for mode in "${modes[@]}"; do + case "${mode}" in + on|off) ;; + *) echo "unsupported mode: ${mode}" >&2; exit 2 ;; + esac +done + +for tool in curl jq; do + command -v "${tool}" >/dev/null || { echo "${tool} is required" >&2; exit 2; } +done + +port_base="${E2E_PORT_BASE:-27100}" +work_dir="${E2E_WORK_DIR:-$(mktemp -d "${TMPDIR:-/tmp}/stargate-e2e-last-cluster.XXXXXX")}" +mkdir -p "${work_dir}" +gateway_bin="${work_dir}/llm-api-gateway" + +model="e2e-model" +routing_key="e2e" +session_id="e2e-last-cluster-session" +follow_up_count=10 +# Each mock-dynamo token takes token_delay_ms. Holders generate +# holder_tokens tokens, so they keep A full for about 3 seconds. +token_delay_ms=50 +holder_tokens=60 +session_tokens=2 +max_engine_concurrency=2 + +sg_grpc=$((port_base)) +sg_http=$((port_base + 1)) +sg_metrics=$((port_base + 2)) +sg_discovery=$((port_base + 3)) +sg_reverse=$((port_base + 4)) +worker_auth=$((port_base + 5)) +gw_http=$((port_base + 30)) +gw_metrics=$((port_base + 31)) +mock_port() { echo $((port_base + 10 + 10 * $1)); } +pylon_metrics_port() { echo $((port_base + 11 + 10 * $1)); } +cluster_names=(a b) + +pids=() +cleanup() { + local pid + for pid in "${pids[@]+"${pids[@]}"}"; do + kill "${pid}" 2>/dev/null || true + done + for pid in "${pids[@]+"${pids[@]}"}"; do + wait "${pid}" 2>/dev/null || true + done + pids=() +} +on_exit() { + local status=$? + cleanup + if [[ ${status} -ne 0 ]]; then + echo "FAIL: logs kept in ${work_dir}" >&2 + elif [[ "${E2E_KEEP_WORK_DIR:-0}" == "1" || -n "${E2E_WORK_DIR:-}" ]]; then + echo "logs kept in ${work_dir}" + else + rm -rf "${work_dir}" + fi + exit "${status}" +} +trap on_exit EXIT +trap 'exit 130' INT TERM + +fail() { + echo "FAIL: $*" >&2 + exit 1 +} + +port_open() { + (exec 3<>"/dev/tcp/127.0.0.1/$1") 2>/dev/null +} + +check_ports_free() { + local port + for port in "$@"; do + if port_open "${port}"; then + fail "port ${port} is already in use; set E2E_PORT_BASE" + fi + done +} + +wait_for() { + local description="$1" deadline=$((SECONDS + 30)) + shift + until "$@" >/dev/null 2>&1; do + [[ ${SECONDS} -lt ${deadline} ]] || fail "timed out waiting for ${description}" + sleep 0.2 + done +} + +build_binaries() { + echo "building stargate, pylon, mock-dynamo, and the worker-auth fixture (release)" + (cd "${stargate_root}" && cargo build --locked --release -q -p stargate -p pylon -p mock-dynamo) + echo "building llm-api-gateway" + (cd "${gateway_root}" && go build -o "${gateway_bin}" ./cmd/llm-api-gateway) +} + +write_lb_config() { + local algorithm="$1" flag="$2" path="$3" virtual_nodes="" + if [[ "${algorithm}" == "wait-and-widen" ]]; then + virtual_nodes='"cache_affinity_virtual_nodes": 8,' + fi + cat >"${path}" <"${log}" 2>&1 & + pids+=("$!") +} + +stargate_metric_sum() { + local name="$1" filter="$2" + curl -sf "http://127.0.0.1:${sg_metrics}/metrics" | + awk -v name="${name}" -v filter="${filter}" ' + index($0, name "{") == 1 && index($0, filter) > 0 { sum += $NF } + END { printf "%d\n", sum }' +} + +mock_count() { + curl -sf "http://127.0.0.1:$(mock_port "$1")/test-control" | + jq --arg model "${model}" '[.counters[] + | select(.endpoint == "chat_completions" and .model == $model and .request_class == "api_gateway") + | .count] | add // 0' +} + +# Prints the cluster whose mock-dynamo served the request run by "$@". +served_by() { + local before_a before_b after_a after_b + before_a=$(mock_count 0) + before_b=$(mock_count 1) + "$@" + after_a=$(mock_count 0) + after_b=$(mock_count 1) + case "$((after_a - before_a)):$((after_b - before_b))" in + 1:0) echo "cluster-a" ;; + 0:1) echo "cluster-b" ;; + *) fail "could not attribute request: a +$((after_a - before_a)), b +$((after_b - before_b))" ;; + esac +} + +affinity_key() { + printf 'mt:v1:session:%s' "$(printf '%s' "${session_id}" | shasum -a 256 | cut -d' ' -f1)" +} + +# Sends a direct Stargate request with the session's affinity key and no +# last-cluster hint. Holders use this to occupy the session's rank-1 cluster. +stargate_request() { + local request_id="$1" tokens="$2" + curl -sf -N -o /dev/null "http://127.0.0.1:${sg_http}/v1/chat/completions" \ + -H 'content-type: application/json' \ + -H "x-request-id: ${request_id}" \ + -H "x-routing-key: ${routing_key}" \ + -H "x-model: ${model}" \ + -H 'x-input-tokens: 8' \ + -H "x-cache-affinity-key: $(affinity_key)" \ + -d "{\"model\":\"${model}\",\"stream\":true,\"max_tokens\":${tokens},\"messages\":[{\"role\":\"user\",\"content\":\"hold\"}]}" +} + +# Sends a gateway request in session S. The gateway derives the affinity key +# from prompt_cache_key and adds the last-cluster hint when enabled. +gateway_request() { + local status + status=$(curl -s -N -o "${run_dir}/last-gateway-response.txt" -w '%{http_code}' \ + "http://127.0.0.1:${gw_http}/v1/chat/completions" \ + -H 'content-type: application/json' \ + -d "{\"model\":\"${routing_key}/${model}\",\"stream\":true,\"max_tokens\":${session_tokens},\"prompt_cache_key\":\"${session_id}\",\"messages\":[{\"role\":\"user\",\"content\":\"turn\"}]}") + [[ "${status}" == "200" ]] || fail "gateway request returned ${status}: $(cat "${run_dir}/last-gateway-response.txt")" +} + +start_stack() { + local algorithm="$1" mode="$2" flag=false gateway_flag=false index name + [[ "${mode}" == "on" ]] && flag=true && gateway_flag=true + check_ports_free "${sg_grpc}" "${sg_http}" "${sg_metrics}" "${sg_discovery}" \ + "${worker_auth}" "${gw_http}" "${gw_metrics}" \ + "$(mock_port 0)" "$(mock_port 1)" "$(pylon_metrics_port 0)" "$(pylon_metrics_port 1)" + write_lb_config "${algorithm}" "${flag}" "${run_dir}/lb-config.json" + + start worker-auth "${run_dir}/worker-auth.log" \ + "${release_dir}/stargate-worker-auth-gateway" \ + --listen-addr "127.0.0.1:${worker_auth}" \ + --worker "token-a=${routing_key}" --worker "token-b=${routing_key}" + wait_for "worker-auth fixture" port_open "${worker_auth}" + start stargate "${run_dir}/stargate.log" env RUST_LOG=info \ + "${release_dir}/stargate" \ + --stargate-id e2e-stargate \ + --listen-addr "127.0.0.1:${sg_grpc}" \ + --http-listen-addr "127.0.0.1:${sg_http}" \ + --model-discovery-listen-addr "127.0.0.1:${sg_discovery}" \ + --advertise-addr "127.0.0.1:${sg_grpc}" \ + --stargate-discovery-dns-name localhost \ + --disable-dns-discovery \ + --metrics-port "${sg_metrics}" \ + --lb-config-path "${run_dir}/lb-config.json" \ + --backend-connectivity=reverse \ + --reverse-tunnel-listen-addr "127.0.0.1:${sg_reverse}" \ + --advertised-hostname-template localhost \ + --worker-auth-endpoint "http://127.0.0.1:${worker_auth}" \ + --quic-insecure + + for index in 0 1; do + name="${cluster_names[${index}]}" + start "mock-${name}" "${run_dir}/mock-${name}.log" \ + "${release_dir}/mock-dynamo" \ + --http-listen-addr "127.0.0.1:$(mock_port "${index}")" \ + --model-name "${model}" \ + --num-tokens "${holder_tokens}" \ + --token-delay-ms "${token_delay_ms}" + done + wait_for "stargate metrics" curl -sf "http://127.0.0.1:${sg_metrics}/metrics" + for index in 0 1; do + name="${cluster_names[${index}]}" + wait_for "mock-dynamo ${name}" curl -sf "http://127.0.0.1:$(mock_port "${index}")/health" + start "pylon-${name}" "${run_dir}/pylon-${name}.log" env RUST_LOG=info \ + "${release_dir}/pylon" \ + --upstream-http-base-url "http://127.0.0.1:$(mock_port "${index}")" \ + --model-name "${model}" \ + --stargate-address "127.0.0.1:${sg_grpc}" \ + --inference-server-id "pylon-${name}" \ + --cluster-id "cluster-${name}" \ + --auth-token "token-${name}" \ + --backend-connectivity reverse \ + --quic-insecure \ + --kv-cache-stats-path /kv-cache/stats \ + --min-update-interval-ms 100 \ + --disable-bringup \ + --active-canary-interval-ms=0 \ + --initial-input-tps 1000 \ + --max-engine-concurrency "${max_engine_concurrency}" \ + --metrics-host 127.0.0.1 \ + --metrics-port "$(pylon_metrics_port "${index}")" + done + + start gateway "${run_dir}/gateway.log" env \ + NVCF_GATEWAY_ADDR="127.0.0.1:${gw_http}" \ + METRICS_PORT="${gw_metrics}" \ + NVCF_REGION=local \ + LOCAL_FUNCTION_ID="${routing_key}" \ + NVCF_DEFAULT_MODEL="${model}" \ + OLRIC_ENABLED=false \ + RATE_LIMIT_ENABLED=false \ + STARGATE_URL="http://127.0.0.1:${sg_http}" \ + STARGATE_LAST_CLUSTER_ENABLED="${gateway_flag}" \ + "${gateway_bin}" + wait_for "gateway readiness" curl -sf "http://127.0.0.1:${gw_http}/readyz" + wait_for "both pylon reverse tunnels" both_tunnels_connected + # Both clusters must be routable. A warm-up through Stargate with a + # throwaway affinity key confirms the routing target exists. + wait_for "stargate routing" curl -sf -N -o /dev/null \ + "http://127.0.0.1:${sg_http}/v1/chat/completions" \ + -H 'content-type: application/json' -H 'x-request-id: e2e-warmup' \ + -H "x-routing-key: ${routing_key}" -H "x-model: ${model}" \ + -H 'x-input-tokens: 1' -H 'x-cache-affinity-key: e2e-warmup' \ + -d "{\"model\":\"${model}\",\"stream\":true,\"max_tokens\":1,\"messages\":[{\"role\":\"user\",\"content\":\"warm\"}]}" +} + +both_tunnels_connected() { + grep -q "reverse tunnel connected" "${run_dir}/pylon-a.log" && + grep -q "reverse tunnel connected" "${run_dir}/pylon-b.log" +} + +run_case() { + local algorithm="$1" mode="$2" + run_dir="${work_dir}/${algorithm}-${mode}" + mkdir -p "${run_dir}" + echo + echo "=== ${algorithm}, last-cluster ${mode} ===" + start_stack "${algorithm}" "${mode}" + # Let both Pylons publish a stats update after the warm-up. + sleep 1 + + local holder_pids=() rank_one rank_two first overflow follow_ups=() + local session_before returning_before new_fallback_before + session_before=$(stargate_metric_sum stargate_routing_session_selections_total "model=\"${model}\"") + + # Fill rank 1. Each holder is a direct Stargate request in session S with + # no hint, so it lands on the session's rank-1 cluster while it has room. + local before_a before_b + before_a=$(mock_count 0) + before_b=$(mock_count 1) + for index in $(seq 1 "${max_engine_concurrency}"); do + stargate_request "e2e-holder-${index}" "${holder_tokens}" & + holder_pids+=("$!") + sleep 0.3 + done + sleep 0.3 + local holders_a=$(($(mock_count 0) - before_a)) holders_b=$(($(mock_count 1) - before_b)) + if [[ ${holders_a} -eq ${max_engine_concurrency} && ${holders_b} -eq 0 ]]; then + rank_one="cluster-a" + rank_two="cluster-b" + elif [[ ${holders_b} -eq ${max_engine_concurrency} && ${holders_a} -eq 0 ]]; then + rank_one="cluster-b" + rank_two="cluster-a" + else + fail "holders did not land on one cluster: a=${holders_a} b=${holders_b}" + fi + echo "session rank 1 (A): ${rank_one}; rank 2 (B): ${rank_two}; A is full" + + returning_before=$(stargate_metric_sum stargate_routing_session_selections_total 'session_state="returning"') + new_fallback_before=$(stargate_metric_sum stargate_routing_session_selections_total 'selection="fallback",session_state="new"') + + first=$(served_by gateway_request) + echo "first gateway request (no hint): ${first}" + [[ "${first}" == "${rank_two}" ]] || fail "first request should overflow to ${rank_two}, got ${first}" + overflow="${first}" + + for pid in "${holder_pids[@]}"; do + wait "${pid}" || fail "holder request failed" + done + # A drained. Give Pylon time to publish the idle stats. + sleep 1 + echo "A drained" + + local index + for index in $(seq 1 "${follow_up_count}"); do + follow_ups+=("$(served_by gateway_request)") + # The gateway stores the last cluster in the background. + sleep 0.1 + done + echo "follow-up clusters: ${follow_ups[*]}" + + local expected on_b=0 on_a=0 cluster + for cluster in "${follow_ups[@]}"; do + if [[ "${cluster}" == "${overflow}" ]]; then on_b=$((on_b + 1)); else on_a=$((on_a + 1)); fi + done + echo "follow-ups on B: ${on_b}/${follow_up_count}; on A: ${on_a}/${follow_up_count}" + + local returning new_fallback session_total + returning=$(($(stargate_metric_sum stargate_routing_session_selections_total 'session_state="returning"') - returning_before)) + new_fallback=$(($(stargate_metric_sum stargate_routing_session_selections_total 'selection="fallback",session_state="new"') - new_fallback_before)) + session_total=$(($(stargate_metric_sum stargate_routing_session_selections_total "model=\"${model}\"") - session_before)) + echo "stargate session selections: returning=${returning} new-fallback=${new_fallback} total=${session_total}" + curl -sf "http://127.0.0.1:${gw_metrics}/metrics" | grep -E '^llm_api_gateway_last_cluster_(lookups|writes)_total' | grep -v ' 0$' || true + + if [[ "${mode}" == "on" ]]; then + expected="${overflow}" + [[ ${on_b} -eq ${follow_up_count} ]] || fail "flag on: expected ${follow_up_count}/${follow_up_count} follow-ups on ${expected}" + [[ ${returning} -eq ${follow_up_count} ]] || fail "flag on: expected ${follow_up_count} returning selections, got ${returning}" + [[ ${new_fallback} -eq 1 ]] || fail "flag on: expected 1 new fallback selection, got ${new_fallback}" + else + expected="${rank_one}" + [[ ${on_a} -eq ${follow_up_count} ]] || fail "flag off: expected ${follow_up_count}/${follow_up_count} follow-ups on ${expected}" + [[ ${session_total} -eq 0 ]] || fail "flag off: expected no session selections, got ${session_total}" + fi + echo "PASS: ${algorithm}, last-cluster ${mode}" + cleanup +} + +if [[ ${build} -eq 1 ]]; then + build_binaries +else + [[ -x "${gateway_bin}" ]] || (cd "${gateway_root}" && go build -o "${gateway_bin}" ./cmd/llm-api-gateway) +fi +for binary in stargate pylon mock-dynamo stargate-worker-auth-gateway; do + [[ -x "${release_dir}/${binary}" ]] || fail "missing ${release_dir}/${binary}; run without --no-build" +done + +for algorithm in "${algorithms[@]}"; do + for mode in "${modes[@]}"; do + run_case "${algorithm}" "${mode}" + done +done +echo +echo "all cases passed" From 7c900f1195e19117ec5dd1432860898cf5652109 Mon Sep 17 00:00:00 2001 From: jcameron Date: Sun, 11 Oct 2026 23:05:11 -0300 Subject: [PATCH 4/6] fix(stargate): pin session classification when the Pulsar ranking is empty pulsar-wait-and-widen returned Unavailable for an empty ranking before it built the session plan, so the first decision left the request unclassified. If weights or cluster membership changed before the next attempt, the request could be classified differently, for example as returning instead of stale, and wait for a hinted cluster it should not wait for. Build the plan first so the first decision always pins the classification, then return Unavailable for an empty ranking. Refs #2289 Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../load_balancer/pulsar_wait_and_widen.rs | 48 +++++++++++++++++-- 1 file changed, 44 insertions(+), 4 deletions(-) diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs index d66b996a77..6c85236906 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs @@ -151,12 +151,12 @@ impl PulsarWaitAndWidenLoadBalancer { elapsed: Duration, ) -> LoadBalancerDecision { let ranking = self.ranking.compute_ranking(request, candidates); - if ranking.is_empty() { - return LoadBalancerDecision::Unavailable; - } // Session rules reorder a per-request copy of the ranking. Rank 1 has // no shortcut: it goes through queue admission, the affinity wait, and - // the affinity discount like every other affinity member. + // the affinity discount like every other affinity member. The plan is + // built before the empty-ranking check so the first decision always + // pins the session classification, even when no candidate has a + // usable Pulsar weight yet. let plan = session::plan_affinity( self.last_cluster_affinity, request, @@ -167,6 +167,9 @@ impl PulsarWaitAndWidenLoadBalancer { self.affinity_input_tokens_scale, ); let ranked_indices = &*plan.order; + if ranked_indices.is_empty() { + return LoadBalancerDecision::Unavailable; + } let affinity_wait = plan.affinity_wait; let band_widen_interval = self.band_widen_interval.unwrap_or(affinity_wait); @@ -1397,6 +1400,43 @@ mod tests { } } + #[test] + fn empty_ranking_still_pins_first_classification() { + let target = target(); + let mut unweighted = single_slot_candidates(1); + unweighted[0].stats.max_input_tps = None; + let hint = LastClusterHint::new(Some("cluster-1".to_string())); + let hybrid = session_hybrid("empty-ranking-seed", true, |_| {}); + let request = session_request(&target, 0, &hint); + + assert_eq!( + hybrid.decide_at(&request, &unweighted, Duration::ZERO), + LoadBalancerDecision::Unavailable + ); + assert_eq!( + hint.classification(), + Some(SessionClassification { + state: SessionState::Stale, + promoted: false, + }) + ); + + // The hinted cluster joins before the retry. The request stays stale: + // no promotion and no affinity wait, so a free cluster wins at once. + let mut candidates = single_slot_candidates(3); + candidates[1].stats.num_running_queries = 1; + let choice = hybrid + .decide_at(&request, &candidates, Duration::ZERO) + .selected() + .expect("stale session must not wait for the hinted cluster"); + assert_ne!(choice.candidate_index, 1); + assert_eq!( + hint.classification() + .map(|classification| classification.state), + Some(SessionState::Stale) + ); + } + #[test] fn last_cluster_needs_an_affinity_key() { let target = target(); From 710ed84f9a5febf3328a43e9702cb7c5ef84cbca Mon Sep 17 00:00:00 2001 From: jcameron Date: Sun, 11 Oct 2026 23:10:55 -0300 Subject: [PATCH 5/6] fix(llm-api-gateway): bound last-cluster store calls and drain after handlers stop The embedded Olric store can hold an operation past its context deadline, so lookup goroutines could outlive the 20ms lookup timeout and write goroutines the 1s write timeout. Under sustained load or a stuck store they could pile up without limit. Cap in-flight lookups and writes at 1024 each, holding a slot until the store call returns. When a cap is reached, the lookup returns at once without the header (result=timeout) and the write is dropped (result=error), with sampled warn logs. The client request is never blocked. Drain could return while a write that started during the wait was still running. It now rechecks the in-flight count after each idle signal. http.Server.Shutdown does not wait for its hooks, so pending writes could start after the drain and reach an Olric node that had already left the cluster. server.New now returns a Gateway whose Shutdown waits for in-flight handlers, then for the rate-limit sync runtime stop (unchanged, still started from the shutdown hook), then drains pending writes, and only then stops the Olric collector and node, each bounded by OLRIC_SHUTDOWN_TIMEOUT. Refs #2289 Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../llm-api-gateway/README.md | 7 + .../cmd/llm-api-gateway/main.go | 2 +- .../llm-api-gateway/lastcluster/tracker.go | 192 +++++++++++++----- .../lastcluster/tracker_test.go | 150 ++++++++++++++ .../server/last_cluster_test.go | 61 ++++++ .../llm-api-gateway/server/server.go | 89 +++++--- .../llm-api-gateway/server/server_test.go | 2 +- 7 files changed, 428 insertions(+), 75 deletions(-) diff --git a/src/invocation-plane-services/llm-api-gateway/README.md b/src/invocation-plane-services/llm-api-gateway/README.md index 8ab3c28507..c28d1f24eb 100644 --- a/src/invocation-plane-services/llm-api-gateway/README.md +++ b/src/invocation-plane-services/llm-api-gateway/README.md @@ -91,6 +91,13 @@ uses the hint only for models whose routing algorithm sets and never delay or fail the client response. - The lookup is bounded by `STARGATE_LAST_CLUSTER_LOOKUP_TIMEOUT`. On a miss, error, or timeout the request proceeds without the hint. +- At most 1024 lookups and 1024 writes may wait on the store at once, counting + calls that outlived their timeout. When a cap is reached, a lookup returns at + once without the hint and counts as `timeout`, and a write is dropped and + counts as `error`. Requests never wait for a slot. +- On shutdown the gateway stops serving, waits for in-flight requests, drains + pending writes, and then stops the Olric node. The drain and the Olric + shutdown are each bounded by `OLRIC_SHUTDOWN_TIMEOUT`. - Store keys are `lc:v1:` plus the hex SHA-256 of the length-prefixed routing key, model, and cache affinity key. Raw routing keys never appear in store keys. Affinity keys and session IDs never appear in store keys or logs. diff --git a/src/invocation-plane-services/llm-api-gateway/cmd/llm-api-gateway/main.go b/src/invocation-plane-services/llm-api-gateway/cmd/llm-api-gateway/main.go index 710efc44f7..5b54c2304f 100644 --- a/src/invocation-plane-services/llm-api-gateway/cmd/llm-api-gateway/main.go +++ b/src/invocation-plane-services/llm-api-gateway/cmd/llm-api-gateway/main.go @@ -89,7 +89,7 @@ func main() { ctx, cfg.Server.Addr, shutdownTimeout(cfg.Server.WriteTimeout), - func(addr string) error { return server.Start(e, addr) }, + func(addr string) error { return server.Start(e.Echo, addr) }, e.Shutdown, ); err != nil { zlog.Fatal().Err(err).Msg("gateway exited unexpectedly") diff --git a/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker.go b/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker.go index b19c966563..ab2ee89a23 100644 --- a/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker.go +++ b/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker.go @@ -36,12 +36,28 @@ import ( // DefaultWriteTimeout bounds one background store write. const DefaultWriteTimeout = time.Second +// Default caps on store calls still running in the background. The Olric +// client does not always honor context cancellation, so a lookup can outlive +// its deadline and a write can outlive its timeout. The caps bound the +// goroutines, and the request contexts they hold, when the store is slow or +// stuck. At the cap, lookups and writes are skipped instead of queued. +const ( + DefaultMaxInflightLookups = 1024 + DefaultMaxInflightWrites = 1024 +) + +var errCapacity = errors.New("last-cluster store has too many calls in flight") + // Options configures a Tracker. type Options struct { TTL time.Duration LookupTimeout time.Duration // WriteTimeout bounds one background write. Zero uses DefaultWriteTimeout. WriteTimeout time.Duration + // MaxInflightLookups and MaxInflightWrites cap store calls still running. + // Zero uses the defaults. + MaxInflightLookups int + MaxInflightWrites int } // Tracker looks up and records the last cluster per session. A nil Tracker is @@ -50,19 +66,35 @@ type Tracker struct { store Store opts Options - // inflight counts background writes. idle is closed when it drops to - // zero. A WaitGroup is not used because writes may start while Drain is - // waiting during shutdown. - mu sync.Mutex - inflight int - idle chan struct{} + // lookupSlots holds one token per lookup whose store call has not + // returned yet. + lookupSlots chan struct{} + + // writes counts background writes. idle is closed when it drops to zero. + // A WaitGroup is not used because writes may start while Drain waits. + mu sync.Mutex + writes int + idle chan struct{} + + // drainWaitHook runs each time Drain is about to wait. Tests only. + drainWaitHook func() } func NewTracker(store Store, opts Options) *Tracker { if opts.WriteTimeout <= 0 { opts.WriteTimeout = DefaultWriteTimeout } - return &Tracker{store: store, opts: opts} + if opts.MaxInflightLookups <= 0 { + opts.MaxInflightLookups = DefaultMaxInflightLookups + } + if opts.MaxInflightWrites <= 0 { + opts.MaxInflightWrites = DefaultMaxInflightWrites + } + return &Tracker{ + store: store, + opts: opts, + lookupSlots: make(chan struct{}, opts.MaxInflightLookups), + } } // eligibleSource reports whether a session source yields a key that stays @@ -113,29 +145,19 @@ func (t *Tracker) Lookup( setFunctionID(span, reqCtx.RoutingKey) lookupCtx = trace.ContextWithSpan(lookupCtx, span) - // The store call runs in its own goroutine so a store that ignores - // context cancellation still cannot hold the request past the deadline. - done := make(chan lookupResult, 1) - go func() { - value, ok, err := t.store.Get(lookupCtx, key) - done <- lookupResult{value: value, ok: ok, err: err} - }() - var res lookupResult select { - case res = <-done: - case <-lookupCtx.Done(): - // Prefer a result that landed at the deadline. - select { - case res = <-done: - default: - res.err = lookupCtx.Err() - } + case t.lookupSlots <- struct{}{}: + res = t.runLookup(lookupCtx, key) + default: + res.err = errCapacity } result := telemetry.LastClusterLookupMiss switch { - case errors.Is(res.err, context.DeadlineExceeded), errors.Is(res.err, context.Canceled): + case errors.Is(res.err, context.DeadlineExceeded), + errors.Is(res.err, context.Canceled), + errors.Is(res.err, errCapacity): result = telemetry.LastClusterLookupTimeout case res.err != nil: result = telemetry.LastClusterLookupError @@ -156,8 +178,7 @@ func (t *Tracker) Lookup( if result == telemetry.LastClusterLookupTimeout { // A slow store times out every lookup at once; the counter and span // already count each one, so sample the log. - sampled := logger.Sample(zerolog.Often) - logger = &sampled + logger = sampled(logger) } logger.Warn(). Err(res.err). @@ -169,6 +190,45 @@ func (t *Tracker) Lookup( return key, clusterID } +// runLookup calls the store in its own goroutine so a store that ignores +// context cancellation still cannot hold the request past the deadline. The +// caller has taken a lookup slot; the goroutine frees it when the store call +// returns, not at the deadline. +func (t *Tracker) runLookup(ctx context.Context, key string) lookupResult { + done := make(chan lookupResult, 1) + go func() { + defer func() { <-t.lookupSlots }() + value, ok, err := t.store.Get(ctx, key) + done <- lookupResult{value: value, ok: ok, err: err} + }() + + select { + case res := <-done: + return res + case <-ctx.Done(): + // Prefer a result that landed at the deadline. + select { + case res := <-done: + return res + default: + return lookupResult{err: ctx.Err()} + } + } +} + +// InflightLookups returns the number of store lookups that have not returned. +func (t *Tracker) InflightLookups() int { + if t == nil { + return 0 + } + return len(t.lookupSlots) +} + +func sampled(logger *zerolog.Logger) *zerolog.Logger { + l := logger.Sample(zerolog.Often) + return &l +} + // setFunctionID tags the span with the function ID, which the gateway takes // from the routing key. func setFunctionID(span trace.Span, routingKey string) { @@ -179,7 +239,8 @@ func setFunctionID(span trace.Span, routingKey string) { // Remember records clusterID for the session in the background when status is // 2xx. key and model are the values Lookup used. It never blocks the caller. -// Empty, oversized, or non-printable cluster IDs are ignored. +// Empty, oversized, or non-printable cluster IDs are ignored. When the write +// cap is reached, the write is dropped and counted as an error. func (t *Tracker) Remember( ctx context.Context, reqCtx *requestctx.RequestContext, @@ -199,11 +260,25 @@ func (t *Tracker) Remember( if reqCtx != nil { requestID, routingKey = reqCtx.RequestID, reqCtx.RoutingKey } + logFailure := func(logger *zerolog.Logger, err error) { + logger.Warn(). + Err(err). + Str("request_id", requestID). + Str("model", model). + Str("routing_key", routingKey). + Msg("stargate last-cluster write failed") + } + + if !t.beginWrite() { + telemetry.RecordLastClusterWrite(ctx, telemetry.LastClusterWriteError) + logFailure(sampled(telemetry.Logger(ctx)), errCapacity) + return + } // The write outlives the request: keep the trace and logger from ctx but - // not its cancellation. + // not its cancellation. The goroutine, and the context it holds, end when + // the store call returns. writeCtx := context.WithoutCancel(ctx) - t.beginWrite() go func() { defer t.endWrite() writeCtx, cancel := context.WithTimeout(writeCtx, t.opts.WriteTimeout) @@ -218,50 +293,67 @@ func (t *Tracker) Remember( result = telemetry.LastClusterWriteError span.RecordError(err) span.SetStatus(codes.Error, "last cluster write failed") - telemetry.Logger(writeCtx).Warn(). - Err(err). - Str("request_id", requestID). - Str("model", model). - Str("routing_key", routingKey). - Msg("stargate last-cluster write failed") + logFailure(telemetry.Logger(writeCtx), err) } span.SetAttributes(attribute.String("result", result)) telemetry.RecordLastClusterWrite(writeCtx, result) }() } -func (t *Tracker) beginWrite() { +// beginWrite reserves a write slot. It returns false at the write cap. +func (t *Tracker) beginWrite() bool { t.mu.Lock() defer t.mu.Unlock() - if t.inflight == 0 { + if t.writes >= t.opts.MaxInflightWrites { + return false + } + if t.writes == 0 { t.idle = make(chan struct{}) } - t.inflight++ + t.writes++ + return true } func (t *Tracker) endWrite() { t.mu.Lock() defer t.mu.Unlock() - t.inflight-- - if t.inflight == 0 { + t.writes-- + if t.writes == 0 { close(t.idle) } } -// Drain waits until no background write is in flight or ctx is done. -func (t *Tracker) Drain(ctx context.Context) { +// InflightWrites returns the number of store writes that have not returned. +func (t *Tracker) InflightWrites() int { if t == nil { - return + return 0 } t.mu.Lock() - if t.inflight == 0 { - t.mu.Unlock() + defer t.mu.Unlock() + return t.writes +} + +// Drain waits until no background write is in flight or ctx is done. Writes +// that start while Drain waits are waited for too. +func (t *Tracker) Drain(ctx context.Context) { + if t == nil { return } - idle := t.idle - t.mu.Unlock() - select { - case <-idle: - case <-ctx.Done(): + for { + t.mu.Lock() + if t.writes == 0 { + t.mu.Unlock() + return + } + idle := t.idle + t.mu.Unlock() + if t.drainWaitHook != nil { + t.drainWaitHook() + } + select { + case <-idle: + case <-ctx.Done(): + return + } } } diff --git a/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker_test.go b/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker_test.go index 79f4c426cf..fee344007c 100644 --- a/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker_test.go +++ b/src/invocation-plane-services/llm-api-gateway/lastcluster/tracker_test.go @@ -22,6 +22,7 @@ import ( "context" "errors" "net/http" + "runtime" "strings" "sync" "testing" @@ -51,6 +52,7 @@ type fakeStore struct { getErr error setErr error blockGet chan struct{} + blockSet chan struct{} } func newFakeStore() *fakeStore { @@ -73,6 +75,12 @@ func (s *fakeStore) Get(_ context.Context, key string) (string, bool, error) { } func (s *fakeStore) Set(_ context.Context, key, value string, _ time.Duration) error { + s.mu.Lock() + block := s.blockSet + s.mu.Unlock() + if block != nil { + <-block + } s.mu.Lock() defer s.mu.Unlock() s.sets++ @@ -521,3 +529,145 @@ func (b *syncBuffer) String() string { defer b.mu.Unlock() return b.buf.String() } + +func TestTrackerCapsStuckStoreCalls(t *testing.T) { + reader := installMetricReader(t) + ctx := context.Background() + store := newFakeStore() + unblock := make(chan struct{}) + store.blockGet = unblock + store.blockSet = unblock + const maxInflight = 4 + const calls = 50 + tracker := NewTracker(store, Options{ + TTL: time.Minute, + LookupTimeout: 5 * time.Millisecond, + MaxInflightLookups: maxInflight, + MaxInflightWrites: maxInflight, + }) + reqCtx := sessionCtx(requestctx.SessionSourceHeader) + baseline := runtime.NumGoroutine() + + start := time.Now() + for i := 0; i < calls; i++ { + key, clusterID := tracker.Lookup(ctx, reqCtx, reqCtx.Model) + if key == "" || clusterID != "" { + t.Fatalf("Lookup() = %q, %q; want key and no cluster", key, clusterID) + } + tracker.Remember(ctx, reqCtx, reqCtx.Model, key, http.StatusOK, "cluster-b") + } + // Only the first maxInflight lookups wait for their deadline. + if elapsed := time.Since(start); elapsed > time.Second { + t.Fatalf("%d calls against a stuck store took %s", calls, elapsed) + } + if got := tracker.InflightLookups(); got != maxInflight { + t.Fatalf("InflightLookups() = %d, want %d", got, maxInflight) + } + if got := tracker.InflightWrites(); got != maxInflight { + t.Fatalf("InflightWrites() = %d, want %d", got, maxInflight) + } + if delta := runtime.NumGoroutine() - baseline; delta > 2*maxInflight+4 { + t.Fatalf("goroutines grew by %d, want at most %d", delta, 2*maxInflight+4) + } + lookups := counterValues(t, reader, "llm_api_gateway_last_cluster_lookups_total") + if lookups["timeout"] != calls { + t.Fatalf("lookups = %v, want timeout=%d", lookups, calls) + } + writes := counterValues(t, reader, "llm_api_gateway_last_cluster_writes_total") + if writes["error"] != calls-maxInflight { + t.Fatalf("writes = %v, want error=%d for dropped writes", writes, calls-maxInflight) + } + + close(unblock) + tracker.Drain(ctx) + deadline := time.Now().Add(2 * time.Second) + for tracker.InflightLookups() != 0 || tracker.InflightWrites() != 0 { + if time.Now().After(deadline) { + t.Fatalf("slots not released: lookups=%d writes=%d", tracker.InflightLookups(), tracker.InflightWrites()) + } + time.Sleep(time.Millisecond) + } + if got := counterValues(t, reader, "llm_api_gateway_last_cluster_writes_total"); got["ok"] != maxInflight { + t.Fatalf("writes = %v, want ok=%d once unblocked", got, maxInflight) + } + // Freed slots serve new calls. + if _, clusterID := tracker.Lookup(ctx, reqCtx, reqCtx.Model); clusterID != "cluster-b" { + t.Fatalf("Lookup() after unblock cluster = %q, want cluster-b", clusterID) + } +} + +func TestDrainWaitsForWritesStartedWhileWaiting(t *testing.T) { + ctx := context.Background() + store := &gatedStore{fakeStore: newFakeStore(), gates: map[string]chan struct{}{}} + tracker := newTestTracker(store) + reqCtx := sessionCtx(requestctx.SessionSourceHeader) + gate1, gate2 := store.gate("k1"), store.gate("k2") + + tracker.Remember(ctx, reqCtx, reqCtx.Model, "k1", http.StatusOK, "cluster-a") + + // The hook stops Drain after it captured the idle channel of W1 and + // before it waits on it. + captured := make(chan struct{}) + resume := make(chan struct{}) + var once sync.Once + tracker.drainWaitHook = func() { + once.Do(func() { + close(captured) + <-resume + }) + } + drained := make(chan struct{}) + go func() { + tracker.Drain(ctx) + close(drained) + }() + <-captured + + // W1 finishes and W2 starts before Drain resumes. + close(gate1) + for tracker.InflightWrites() != 0 { + time.Sleep(time.Millisecond) + } + tracker.Remember(ctx, reqCtx, reqCtx.Model, "k2", http.StatusOK, "cluster-b") + close(resume) + + select { + case <-drained: + t.Fatal("Drain returned while W2 was still in flight") + case <-time.After(50 * time.Millisecond): + } + close(gate2) + select { + case <-drained: + case <-time.After(2 * time.Second): + t.Fatal("Drain did not return after W2 finished") + } + if got := store.value("k2"); got != "cluster-b" { + t.Fatalf("stored k2 = %q, want cluster-b", got) + } +} + +// gatedStore blocks Set for each key until its gate is closed. +type gatedStore struct { + *fakeStore + gateMu sync.Mutex + gates map[string]chan struct{} +} + +func (s *gatedStore) gate(key string) chan struct{} { + s.gateMu.Lock() + defer s.gateMu.Unlock() + ch := make(chan struct{}) + s.gates[key] = ch + return ch +} + +func (s *gatedStore) Set(ctx context.Context, key, value string, ttl time.Duration) error { + s.gateMu.Lock() + ch := s.gates[key] + s.gateMu.Unlock() + if ch != nil { + <-ch + } + return s.fakeStore.Set(ctx, key, value, ttl) +} diff --git a/src/invocation-plane-services/llm-api-gateway/server/last_cluster_test.go b/src/invocation-plane-services/llm-api-gateway/server/last_cluster_test.go index 6f85af7ee5..03242ce5d6 100644 --- a/src/invocation-plane-services/llm-api-gateway/server/last_cluster_test.go +++ b/src/invocation-plane-services/llm-api-gateway/server/last_cluster_test.go @@ -18,6 +18,7 @@ limitations under the License. package server import ( + "context" "io" "net/http" "net/http/httptest" @@ -27,6 +28,7 @@ import ( "time" "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/internal/olrictest" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/provider" ) // hintingStargate answers like Stargate with x-stargate-cluster-id set and @@ -175,3 +177,62 @@ func TestStargateClusterIDStrippedFromClientResponses(t *testing.T) { } } } + +func TestStateShutdownOrder(t *testing.T) { + var ( + mu sync.Mutex + order []string + ) + record := func(step string) { + mu.Lock() + defer mu.Unlock() + order = append(order, step) + } + limiterStopped := make(chan struct{}) + go func() { + time.Sleep(20 * time.Millisecond) + record("limiter") + close(limiterStopped) + }() + state := stateShutdown{ + limiterStopped: limiterStopped, + drain: func(ctx context.Context) { + if _, ok := ctx.Deadline(); !ok { + t.Error("drain context has no deadline") + } + record("drain") + }, + stopOlric: func(context.Context) { record("olric") }, + } + + state.run() + + if got := strings.Join(order, ","); got != "limiter,drain,olric" { + t.Fatalf("shutdown order = %s, want limiter,drain,olric", got) + } +} + +func TestGatewayShutdownReleasesStateOnceAfterHTTPShutdown(t *testing.T) { + upstream := newHintingStargate(t, "cluster-b") + cfg := testConfig(upstream.URL) + cfg.Stargate.LastCluster.Enabled = true + inferenceProvider, err := provider.NewStargateProvider(cfg.Stargate) + if err != nil { + t.Fatal(err) + } + g, err := New(cfg, inferenceProvider, nil) + if err != nil { + t.Fatal(err) + } + runs := 0 + g.state.stopOlric = func(context.Context) { runs++ } + + for i := 0; i < 2; i++ { + if err := g.Shutdown(context.Background()); err != nil { + t.Fatalf("Shutdown() error = %v", err) + } + } + if runs != 1 { + t.Fatalf("state released %d times, want 1", runs) + } +} diff --git a/src/invocation-plane-services/llm-api-gateway/server/server.go b/src/invocation-plane-services/llm-api-gateway/server/server.go index 658883a479..e5c5341ac2 100644 --- a/src/invocation-plane-services/llm-api-gateway/server/server.go +++ b/src/invocation-plane-services/llm-api-gateway/server/server.go @@ -21,6 +21,7 @@ import ( "context" "fmt" "io" + "sync" "time" echo "github.com/labstack/echo/v4" @@ -36,11 +37,60 @@ import ( "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/util" ) +// Gateway is the gateway's Echo server plus the shared state that must +// outlive in-flight handlers. Shut it down with Shutdown, not Echo.Shutdown. +type Gateway struct { + *echo.Echo + state stateShutdown + stateOnce sync.Once +} + +// Shutdown stops the HTTP server and waits for in-flight handlers, then +// releases the shared state: pending last-cluster writes are drained and the +// Olric node leaves the cluster. Both steps are bounded by +// OLRIC_SHUTDOWN_TIMEOUT. +func (g *Gateway) Shutdown(ctx context.Context) error { + err := g.Echo.Shutdown(ctx) + g.stateOnce.Do(g.state.run) + return err +} + +// stateShutdown releases shared state after the HTTP server has stopped. +type stateShutdown struct { + // limiterStopped is closed by the server's shutdown hook once the + // rate-limit sync runtime has stopped. Nil when there is no hook. + limiterStopped <-chan struct{} + drain func(context.Context) + stopOlric func(context.Context) + timeout time.Duration +} + +// run drains last-cluster writes and then stops Olric. Olric stops only +// after the rate limiter has stopped and the drain has finished, so neither +// writes to a node that already left the cluster. +func (s stateShutdown) run() { + if s.limiterStopped != nil { + <-s.limiterStopped + } + if s.drain != nil { + timeout := s.timeout + if timeout <= 0 { + timeout = util.DefaultOlricShutdownTimeout + } + ctx, cancel := context.WithTimeout(context.Background(), timeout) + s.drain(ctx) + cancel() + } + if s.stopOlric != nil { + s.stopOlric(context.Background()) + } +} + func New( cfg *config.Config, inferenceProvider provider.InferenceProvider, authClient api.InvocationAuthClient, -) (*echo.Echo, error) { +) (*Gateway, error) { if cfg != nil && cfg.Telemetry.ServiceName != "" { telemetry.SetServiceName(cfg.Telemetry.ServiceName) } @@ -91,20 +141,22 @@ func New( aware.SetLastClusterTracker(tracker) } - if olric.node != nil || tracker != nil { - // One hook keeps the order fixed: stop the limiter, wait briefly for - // pending last-cluster writes, then leave the Olric cluster. This is - // best effort: http.Server.Shutdown does not wait for its hooks, so - // writes still pending when the process exits are dropped. A lost - // hint only costs that session one affinity miss. + state := stateShutdown{ + stopOlric: olric.stop, + timeout: cfg.Olric.ShutdownTimeout, + } + if tracker != nil { + state.drain = tracker.Drain + } + if olric.node != nil { + // The rate-limit sync runtime still stops when shutdown begins, as + // before. The Olric node and pending last-cluster writes outlive the + // in-flight handlers and are released by Gateway.Shutdown. + limiterStopped := make(chan struct{}) + state.limiterStopped = limiterStopped e.Server.RegisterOnShutdown(func() { + defer close(limiterStopped) stopLimiter() - if tracker != nil { - drainCtx, cancel := contextWithOptionalTimeout(cfg.Olric.ShutdownTimeout) - tracker.Drain(drainCtx) - cancel() - } - olric.stop(context.Background()) }) } @@ -115,7 +167,7 @@ func New( ) api.RegisterRoutes(e, handlers) - return e, nil + return &Gateway{Echo: e, state: state}, nil } // Start serves e on addr until the server is shut down. Use it instead of @@ -260,15 +312,6 @@ func newLastClusterTracker( }), nil } -// contextWithOptionalTimeout returns a context without a deadline when -// timeout <= 0. -func contextWithOptionalTimeout(timeout time.Duration) (context.Context, context.CancelFunc) { - if timeout <= 0 { - return context.WithCancel(context.Background()) - } - return context.WithTimeout(context.Background(), timeout) -} - // timeAfterOrForever returns a channel that never fires when timeout <= 0, // and a time.After channel otherwise. It lets the shutdown hook stay in a // single select regardless of whether the user configured a timeout. diff --git a/src/invocation-plane-services/llm-api-gateway/server/server_test.go b/src/invocation-plane-services/llm-api-gateway/server/server_test.go index 11e4eff426..2c8701377e 100644 --- a/src/invocation-plane-services/llm-api-gateway/server/server_test.go +++ b/src/invocation-plane-services/llm-api-gateway/server/server_test.go @@ -326,7 +326,7 @@ func TestNewWrapsEchoWithFinalWriteDeadline(t *testing.T) { } // e.Start would reset the handler to e and drop the final write // deadline; main serves through Start, which keeps it. - if e.Server.Handler == http.Handler(e) { + if e.Server.Handler == http.Handler(e.Echo) { t.Fatal("server handler is the bare Echo instance") } } From b53a848372f4b6cfa85baf4b5f7302b8ea36976d Mon Sep 17 00:00:00 2001 From: jcameron Date: Sun, 11 Oct 2026 23:30:01 -0300 Subject: [PATCH 6/6] test(stargate): let the last-cluster e2e check run without the gateway source With --no-build, the script still resolved the gateway source directory and rebuilt the gateway into the work directory, so a copy of only the Stargate tree could not run it. Resolve the gateway source only when building, and accept a prebuilt gateway through E2E_GATEWAY_BIN. Refs #2289 Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- src/libraries/rust/stargate/scripts/README.md | 10 ++++++---- .../rust/stargate/scripts/e2e-last-cluster.sh | 19 +++++++++++++------ 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/src/libraries/rust/stargate/scripts/README.md b/src/libraries/rust/stargate/scripts/README.md index 0cc03662e3..ef15521ca4 100644 --- a/src/libraries/rust/stargate/scripts/README.md +++ b/src/libraries/rust/stargate/scripts/README.md @@ -76,8 +76,10 @@ scripts/e2e-last-cluster.sh --algorithm pulsar-wait-and-widen --mode on Mode `on` sets `STARGATE_LAST_CLUSTER_ENABLED=true` in the gateway and `last_cluster_affinity: true` in Stargate. Mode `off` disables both. The script builds release binaries and the gateway first unless `--no-build` is -set. It attributes each request to a cluster through the mock-dynamo -`/test-control` counters, checks +set. With `--no-build`, set `E2E_GATEWAY_BIN` to a prebuilt gateway to run +without the gateway source or `go`. It attributes each request to a cluster +through the mock-dynamo `/test-control` counters, checks `stargate_routing_session_selections_total`, and exits non-zero on any -mismatch. It needs `cargo`, `go`, `curl`, and `jq`, uses 32 loopback ports -starting at `E2E_PORT_BASE` (default `27100`), and keeps logs on failure. +mismatch. It needs `curl` and `jq` (plus `cargo` and `go` to build), uses 32 +loopback ports starting at `E2E_PORT_BASE` (default `27100`), and keeps logs +on failure. diff --git a/src/libraries/rust/stargate/scripts/e2e-last-cluster.sh b/src/libraries/rust/stargate/scripts/e2e-last-cluster.sh index fafec82e03..a01f98deae 100755 --- a/src/libraries/rust/stargate/scripts/e2e-last-cluster.sh +++ b/src/libraries/rust/stargate/scripts/e2e-last-cluster.sh @@ -46,13 +46,15 @@ Environment: E2E_PORT_BASE First loopback port, default 27100. Uses 32 ports. E2E_WORK_DIR Logs and configs, default a new mktemp directory. E2E_KEEP_WORK_DIR Set to 1 to keep the work directory after success. + E2E_GATEWAY_BIN Prebuilt llm-api-gateway binary to use with --no-build. + Without it, --no-build still builds the gateway. -Requires cargo, go, curl, and jq. +Requires curl and jq. Building also needs cargo and go. EOF } stargate_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -gateway_root="$(cd "${stargate_root}/../../../invocation-plane-services/llm-api-gateway" && pwd)" +gateway_root="${stargate_root}/../../../invocation-plane-services/llm-api-gateway" release_dir="${stargate_root}/target/release" algorithms=() @@ -89,7 +91,7 @@ done port_base="${E2E_PORT_BASE:-27100}" work_dir="${E2E_WORK_DIR:-$(mktemp -d "${TMPDIR:-/tmp}/stargate-e2e-last-cluster.XXXXXX")}" mkdir -p "${work_dir}" -gateway_bin="${work_dir}/llm-api-gateway" +gateway_bin="${E2E_GATEWAY_BIN:-${work_dir}/llm-api-gateway}" model="e2e-model" routing_key="e2e" @@ -167,11 +169,16 @@ wait_for() { done } +build_gateway() { + [[ -d "${gateway_root}" ]] || fail "missing gateway source at ${gateway_root}; set E2E_GATEWAY_BIN" + echo "building llm-api-gateway" + (cd "${gateway_root}" && go build -o "${gateway_bin}" ./cmd/llm-api-gateway) +} + build_binaries() { echo "building stargate, pylon, mock-dynamo, and the worker-auth fixture (release)" (cd "${stargate_root}" && cargo build --locked --release -q -p stargate -p pylon -p mock-dynamo) - echo "building llm-api-gateway" - (cd "${gateway_root}" && go build -o "${gateway_bin}" ./cmd/llm-api-gateway) + build_gateway } write_lb_config() { @@ -445,7 +452,7 @@ run_case() { if [[ ${build} -eq 1 ]]; then build_binaries else - [[ -x "${gateway_bin}" ]] || (cd "${gateway_root}" && go build -o "${gateway_bin}" ./cmd/llm-api-gateway) + [[ -x "${gateway_bin}" ]] || build_gateway fi for binary in stargate pylon mock-dynamo stargate-worker-auth-gateway; do [[ -x "${release_dir}/${binary}" ]] || fail "missing ${release_dir}/${binary}; run without --no-build"