diff --git a/lib/global-frontend/src/relay.rs b/lib/global-frontend/src/relay.rs index 22241de5ce18..e5a4f1683464 100644 --- a/lib/global-frontend/src/relay.rs +++ b/lib/global-frontend/src/relay.rs @@ -3,14 +3,10 @@ //! Pool load from each pool's PoolRelay. //! -//! Two tasks per pool follow its Relay: one subscribes to the serving-load -//! stream and records the pool's requests in flight and how many it can run -//! at once; the other subscribes to the KV-load stream and records the KV -//! blocks its workers hold. Both go to [`PoolSignals`], where each report -//! stays until it ages out, so a reconnect or a briefly incomplete view does -//! not make a pool's load unknown. A pool's Relay comes from the catalog; -//! when it changes or the pool leaves, the subscription is replaced or ends -//! and the pool's load is dropped. +//! One task per pool follows its Relay's serving-load and KV-load streams over +//! the same connection. Reports stay in [`PoolSignals`] until they age out. +//! The Relay address comes from the catalog; when it changes or the pool +//! leaves, the subscription ends and the pool's load is dropped. use std::collections::HashMap; use std::convert::Infallible; @@ -19,12 +15,13 @@ use std::time::{Duration, Instant}; use anyhow::{Context, Result, anyhow, bail}; use dynamo_global_common::PoolId; -use dynamo_global_router::{LoadReport, PoolSignals}; +use dynamo_global_router::{LoadReport, PoolSignals, RelaySource}; use dynamo_llm::pool_relay::wan::grpc::protocol::{ KvEventRelayClient, KvPoolLoadUpdate, RELAY_CONTRACT_MARKER, ServingLoadUpdate, - SubscribeKvPoolLoadRequest, SubscribeServingLoadRequest, + SubscribeKvPoolLoadRequest, SubscribeServingLoadRequest, validate_protocol_envelope, }; use tokio::sync::watch; +use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use tonic::transport::{Channel, Endpoint}; use url::Url; @@ -39,7 +36,6 @@ const FIRST_RETRY: Duration = Duration::from_millis(500); const LAST_RETRY: Duration = Duration::from_secs(30); /// A stream that lasted this long resets the retry delay. const HEALTHY_STREAM: Duration = Duration::from_secs(30); - /// Keeps the Relay subscriptions of every pool in `pools` that names a /// Relay, until `cancel` fires. pub async fn follow_relays( @@ -47,7 +43,8 @@ pub async fn follow_relays( signals: Arc, cancel: CancellationToken, ) { - let mut subscriptions = HashMap::::new(); + let mut subscriptions = + HashMap::)>::new(); loop { let snapshot = pools.borrow_and_update().clone(); let wanted = snapshot @@ -55,29 +52,34 @@ pub async fn follow_relays( .iter() .filter_map(|pool| Some((pool.id.clone(), (pool.relay.clone()?, pool.model.clone())))) .collect::>(); - subscriptions.retain(|pool, (relay, subscription)| { - let keep = wanted.get(pool).is_some_and(|(want, _)| want == relay); - if !keep { - subscription.cancel(); - signals.forget(pool); - } - keep - }); + let removed = subscriptions + .iter() + .filter(|(pool, (relay, model, _, _))| { + !wanted.get(*pool).is_some_and(|(want_relay, want_model)| { + want_relay == relay && want_model == model + }) + }) + .map(|(pool, _)| pool.clone()) + .collect::>(); + for pool in removed { + let (_, _, subscription, task) = subscriptions.remove(&pool).unwrap(); + subscription.cancel(); + let _ = task.await; + signals.forget(&pool); + } for (pool, (relay, model)) in wanted { if subscriptions.contains_key(&pool) { continue; } let subscription = cancel.child_token(); - for feed in [Feed::Load { model }, Feed::Kv] { - tokio::spawn(subscribe( - pool.clone(), - feed, - relay.clone(), - signals.clone(), - subscription.clone(), - )); - } - subscriptions.insert(pool, (relay, subscription)); + let task = tokio::spawn(subscribe( + pool.clone(), + model.clone(), + relay.clone(), + signals.clone(), + subscription.clone(), + )); + subscriptions.insert(pool, (relay, model, subscription, task)); } tokio::select! { _ = cancel.cancelled() => return, @@ -86,18 +88,9 @@ pub async fn follow_relays( } } -/// Which of a Relay's streams a task follows. -enum Feed { - /// Serving load, for the catalog's `model`. - Load { model: String }, - /// The KV blocks the pool's workers hold. - Kv, -} - -/// Streams one pool's `feed`, reconnecting with a growing delay. async fn subscribe( pool: PoolId, - feed: Feed, + model: String, relay: Url, signals: Arc, cancel: CancellationToken, @@ -105,12 +98,7 @@ async fn subscribe( let mut retry = FIRST_RETRY; loop { let started = Instant::now(); - let stream = async { - match &feed { - Feed::Load { model } => stream_load(&pool, model, &relay, &signals).await, - Feed::Kv => stream_kv(&pool, &relay, &signals).await, - } - }; + let stream = stream_pool(&pool, &model, &relay, &signals); // Stop as soon as the subscription is cancelled, even while // connecting, so a replaced subscription never records over the next. let Err(error) = tokio::select! { @@ -129,63 +117,88 @@ async fn subscribe( } } -/// Records each load update until the stream fails. An incomplete update is -/// skipped, leaving the last report to age out. -async fn stream_load( +/// Both RPCs share one HTTP/2 connection, so a Service can select one Relay +/// Pod for the pair. A failed stream ends the pair and triggers a reconnect. +async fn stream_pool( pool: &PoolId, model: &str, relay: &Url, signals: &PoolSignals, ) -> Result { - let mut client = KvEventRelayClient::new(connect(relay).await?); - let request = SubscribeServingLoadRequest { + let channel = connect(relay).await?; + let mut load_client = KvEventRelayClient::new(channel.clone()); + let mut kv_client = KvEventRelayClient::new(channel); + let load_request = SubscribeServingLoadRequest { subscriber_id: SUBSCRIBER_ID.to_owned(), contract_marker: RELAY_CONTRACT_MARKER, }; - let mut updates = - tokio::time::timeout(SUBSCRIBE_TIMEOUT, client.subscribe_serving_load(request)) - .await - .map_err(|_| anyhow!("subscribe timed out"))? - .context("subscribe")? - .into_inner(); - tracing::info!(%pool, %relay, "following the pool's Relay load"); + let kv_request = SubscribeKvPoolLoadRequest { + subscriber_id: SUBSCRIBER_ID.to_owned(), + contract_marker: RELAY_CONTRACT_MARKER, + }; + let (load, kv) = tokio::time::timeout(SUBSCRIBE_TIMEOUT, async { + tokio::try_join!( + load_client.subscribe_serving_load(load_request), + kv_client.subscribe_kv_pool_load(kv_request), + ) + }) + .await + .map_err(|_| anyhow!("subscribe timed out"))? + .context("subscribe")?; + let mut load_updates = load.into_inner(); + let mut kv_updates = kv.into_inner(); + tracing::info!(%pool, %relay, "following the pool's Relay load and KV load"); loop { - let Some(update) = updates.message().await.context("read")? else { - bail!("the Relay closed the stream"); - }; - match pool_report(&update, model) { - Some(report) => signals.record_load(pool, report), - None => tracing::debug!(%pool, "the Relay's load is incomplete"), + tokio::select! { + update = load_updates.message() => { + let Some(update) = update.context("read serving load")? else { + bail!("the Relay closed the serving-load stream"); + }; + record_serving_update(pool, model, signals, &update)?; + } + update = kv_updates.message() => { + let Some(update) = update.context("read KV load")? else { + bail!("the Relay closed the KV-load stream"); + }; + record_kv_update(pool, signals, &update)?; + } } } } -/// Records each KV-load update until the stream fails. An incomplete update -/// is skipped, leaving the last report to age out. -async fn stream_kv(pool: &PoolId, relay: &Url, signals: &PoolSignals) -> Result { - let mut client = KvEventRelayClient::new(connect(relay).await?); - let request = SubscribeKvPoolLoadRequest { - subscriber_id: SUBSCRIBER_ID.to_owned(), - contract_marker: RELAY_CONTRACT_MARKER, +fn record_serving_update( + pool: &PoolId, + model: &str, + signals: &PoolSignals, + update: &ServingLoadUpdate, +) -> Result<()> { + validate_protocol_envelope(update.protocol_version, update.contract_marker)?; + let relay = update.relay.as_ref().context("missing Relay identity")?; + let source = RelaySource { + runtime_id: relay.drt_instance_id, + incarnation: relay.relay_incarnation, }; - let mut updates = - tokio::time::timeout(SUBSCRIBE_TIMEOUT, client.subscribe_kv_pool_load(request)) - .await - .map_err(|_| anyhow!("subscribe timed out"))? - .context("subscribe")? - .into_inner(); - tracing::info!(%pool, %relay, "following the pool's Relay KV load"); - loop { - let Some(update) = updates.message().await.context("read")? else { - bail!("the Relay closed the stream"); - }; - match pool_kv(&update) { - Some((used, total)) => { - signals.record_kv(pool, used, total, report_window(update.window_ms)) - } - None => tracing::debug!(%pool, "the Relay's KV load is incomplete"), + match pool_report(update, model) { + Some(report) => signals.record_load(pool, source, report), + None => tracing::debug!(%pool, "the Relay's load is incomplete"), + } + Ok(()) +} + +fn record_kv_update(pool: &PoolId, signals: &PoolSignals, update: &KvPoolLoadUpdate) -> Result<()> { + validate_protocol_envelope(update.protocol_version, update.contract_marker)?; + let relay = update.relay.as_ref().context("missing Relay identity")?; + let source = RelaySource { + runtime_id: relay.drt_instance_id, + incarnation: relay.relay_incarnation, + }; + match pool_kv(update) { + Some((used, total)) => { + signals.record_kv(pool, source, used, total, report_window(update.window_ms)) } + None => tracing::debug!(%pool, "the Relay's KV load is incomplete"), } + Ok(()) } fn report_window(window_ms: u64) -> Duration { @@ -272,10 +285,20 @@ fn pool_report(update: &ServingLoadUpdate, model: &str) -> Option { #[cfg(test)] mod tests { + use std::net::SocketAddr; + use std::pin::Pin; + use std::sync::Mutex; + + use dynamo_global_common::{PoolIdDeriver, PoolKey, V1PoolIdDeriver}; + use dynamo_global_router::InFlightTracker; + use dynamo_llm::pool_relay::wan::grpc::protocol as proto; use dynamo_llm::pool_relay::wan::grpc::protocol::{ KvPoolLoadEntry, LoadView, ModelServingLoad, PoolDeploymentStatus, PoolServingLoad, - RequestLifecycleStats, + RELAY_PROTOCOL_VERSION, RelayIdentity, RequestLifecycleStats, }; + use futures::{Stream, StreamExt}; + use tonic::transport::Server; + use tonic::{Request, Response, Status}; use super::*; @@ -314,6 +337,82 @@ mod tests { } } + fn pool_id() -> PoolId { + V1PoolIdDeriver.derive(&PoolKey::new("ohio", "dynamo", "blue").unwrap()) + } + + fn relay_identity(runtime_id: u64) -> RelayIdentity { + RelayIdentity { + drt_instance_id: runtime_id, + relay_incarnation: 1, + } + } + + fn serving_update(runtime_id: u64) -> ServingLoadUpdate { + ServingLoadUpdate { + protocol_version: RELAY_PROTOCOL_VERSION, + contract_marker: RELAY_CONTRACT_MARKER, + relay: Some(relay_identity(runtime_id)), + window_sequence: 1, + ..update(Some(128), &[("qwen", Some(3), Some(20))]) + } + } + + fn kv_update(runtime_id: u64) -> KvPoolLoadUpdate { + KvPoolLoadUpdate { + protocol_version: RELAY_PROTOCOL_VERSION, + contract_marker: RELAY_CONTRACT_MARKER, + relay: Some(relay_identity(runtime_id)), + window_sequence: 1, + window_ms: 2000, + pools: vec![KvPoolLoadEntry { + kv_used_blocks: 400, + total_kv_blocks: 1000, + kv_observed_ranks: 1, + kv_expected_ranks: 1, + ..Default::default() + }], + ..Default::default() + } + } + + #[test] + fn rejects_invalid_load_envelopes_before_recording_state() { + let pool = pool_id(); + let signals = PoolSignals::new(Arc::new(InFlightTracker::new())); + let mut load = serving_update(1); + load.contract_marker = 0; + assert!(record_serving_update(&pool, "qwen", &signals, &load).is_err()); + load = serving_update(1); + load.protocol_version += 1; + assert!(record_serving_update(&pool, "qwen", &signals, &load).is_err()); + assert_eq!(signals.occupancy(&pool), None); + + record_serving_update(&pool, "qwen", &signals, &serving_update(1)).unwrap(); + let mut kv = kv_update(1); + kv.contract_marker = 0; + assert!(record_kv_update(&pool, &signals, &kv).is_err()); + kv = kv_update(1); + kv.protocol_version += 1; + assert!(record_kv_update(&pool, &signals, &kv).is_err()); + assert_eq!(signals.occupancy(&pool).unwrap().kv, None); + + record_kv_update(&pool, &signals, &kv_update(1)).unwrap(); + assert!(signals.occupancy(&pool).unwrap().kv.is_some()); + } + + #[test] + fn keeps_separate_streams_on_the_same_relay_replica() { + let pool = pool_id(); + let signals = PoolSignals::new(Arc::new(InFlightTracker::new())); + record_serving_update(&pool, "qwen", &signals, &serving_update(1)).unwrap(); + record_kv_update(&pool, &signals, &kv_update(2)).unwrap(); + assert_eq!(signals.occupancy(&pool).unwrap().kv, None); + + record_kv_update(&pool, &signals, &kv_update(1)).unwrap(); + assert!(signals.occupancy(&pool).unwrap().kv.is_some()); + } + #[test] fn counts_waiting_and_generating_requests_against_capacity() { let report = pool_report(&update(Some(128), &[("qwen", Some(3), Some(20))]), "qwen"); @@ -372,4 +471,149 @@ mod tests { assert_eq!(pool_kv(&update(vec![entry(0, 0, 4, 4)])), None); assert_eq!(pool_kv(&update(vec![])), None); } + + type TestStream = Pin> + Send + 'static>>; + + struct MockRelay { + serving: ServingLoadUpdate, + kv: KvPoolLoadUpdate, + peers: Arc>>, + } + + fn open_stream(value: T) -> TestStream { + Box::pin( + futures::stream::once(async move { Ok(value) }) + .chain(futures::stream::pending::>()), + ) + } + + #[tonic::async_trait] + impl proto::KvEventRelay for MockRelay { + type WatchKvPoolCatalogStream = TestStream; + type SubscribeKvPoolStream = TestStream; + type SubscribeServingReadinessStream = TestStream; + type SubscribeKvPoolLoadStream = TestStream; + type SubscribeServingLoadStream = TestStream; + + async fn get_relay_info( + &self, + _: Request, + ) -> Result, Status> { + Err(Status::unimplemented("unused in this test")) + } + + async fn watch_kv_pool_catalog( + &self, + _: Request, + ) -> Result, Status> { + Err(Status::unimplemented("unused in this test")) + } + + async fn subscribe_kv_pool( + &self, + _: Request, + ) -> Result, Status> { + Err(Status::unimplemented("unused in this test")) + } + + async fn subscribe_serving_readiness( + &self, + _: Request, + ) -> Result, Status> { + Err(Status::unimplemented("unused in this test")) + } + + async fn subscribe_serving_load( + &self, + request: Request, + ) -> Result, Status> { + self.peers + .lock() + .unwrap() + .push(request.remote_addr().unwrap()); + Ok(Response::new(open_stream(self.serving.clone()))) + } + + async fn subscribe_kv_pool_load( + &self, + request: Request, + ) -> Result, Status> { + self.peers + .lock() + .unwrap() + .push(request.remote_addr().unwrap()); + Ok(Response::new(open_stream(self.kv.clone()))) + } + } + + #[tokio::test] + async fn grpc_load_streams_share_a_connection_and_record_matching_replica() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let incoming = futures::stream::unfold(listener, |listener| async move { + listener + .accept() + .await + .ok() + .map(|(stream, _)| (Ok::<_, std::io::Error>(stream), listener)) + }); + let peers = Arc::new(Mutex::new(Vec::new())); + let server_cancel = CancellationToken::new(); + let server = { + let cancel = server_cancel.clone(); + let relay = MockRelay { + serving: serving_update(7), + kv: kv_update(7), + peers: peers.clone(), + }; + tokio::spawn(async move { + Server::builder() + .add_service(proto::KvEventRelayServer::new(relay)) + .serve_with_incoming_shutdown(incoming, cancel.cancelled()) + .await + .unwrap(); + }) + }; + + let pool = pool_id(); + let signals = Arc::new(PoolSignals::new(Arc::new(InFlightTracker::new()))); + let url = Url::parse(&format!("http://{address}")).unwrap(); + let client = { + let pool = pool.clone(); + let signals = signals.clone(); + tokio::spawn(async move { stream_pool(&pool, "qwen", &url, &signals).await }) + }; + let occupancy = tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Some(occupancy) = signals.occupancy(&pool) + && occupancy.kv.is_some() + { + break occupancy; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert_eq!(occupancy.slots, 23.0 / 128.0); + assert_eq!(occupancy.kv, Some(0.46)); + assert_eq!(occupancy.load, 0.46); + { + let peers = peers.lock().unwrap(); + assert_eq!(peers.len(), 2); + assert_eq!(peers[0], peers[1]); + } + + client.abort(); + let _ = client.await; + server_cancel.cancel(); + let mut server = server; + match tokio::time::timeout(Duration::from_secs(5), &mut server).await { + Ok(result) => result.unwrap(), + Err(_) => { + server.abort(); + panic!("mock Relay server did not stop"); + } + } + } } diff --git a/lib/global-router/src/affinity.rs b/lib/global-router/src/affinity.rs index 074a1e4a8f25..0e5eb85226e7 100644 --- a/lib/global-router/src/affinity.rs +++ b/lib/global-router/src/affinity.rs @@ -184,7 +184,7 @@ mod tests { use crate::load::InFlightTracker; use crate::selector::tests::regional_pool; - use crate::signals::LoadReport; + use crate::signals::{LoadReport, RelaySource}; use super::*; @@ -214,6 +214,10 @@ mod tests { /// bind. fn set_occupancy(pools: &[Pool], signals: &PoolSignals, loads: [f64; 4]) { let window = Duration::from_secs(20); + let source = RelaySource { + runtime_id: 1, + incarnation: 1, + }; for (pool, load) in pools.iter().zip(loads) { let in_flight = (load * 128.0).round() as u64; let report = LoadReport { @@ -222,8 +226,8 @@ mod tests { max_concurrency: Some(128), window, }; - signals.record_load(&pool.id, report); - signals.record_kv(&pool.id, 0, 16384, window); + signals.record_load(&pool.id, source, report); + signals.record_kv(&pool.id, source, 0, 16384, window); } } diff --git a/lib/global-router/src/lib.rs b/lib/global-router/src/lib.rs index e84608ede873..84160019b09c 100644 --- a/lib/global-router/src/lib.rs +++ b/lib/global-router/src/lib.rs @@ -21,4 +21,4 @@ pub use pool::Pool; pub use selector::{ LeastInFlightSelector, PoolSelector, PoolTarget, RoundRobinSelector, RoutingRequest, }; -pub use signals::{LoadReport, PoolOccupancy, PoolSignals}; +pub use signals::{LoadReport, PoolOccupancy, PoolSignals, RelaySource}; diff --git a/lib/global-router/src/signals.rs b/lib/global-router/src/signals.rs index e5f476098404..e094e292958d 100644 --- a/lib/global-router/src/signals.rs +++ b/lib/global-router/src/signals.rs @@ -51,15 +51,28 @@ pub struct PoolOccupancy { /// request reads them. #[derive(Debug)] pub struct PoolSignals { - loads: DashMap, - kv: DashMap, + reports: DashMap>, regions: RwLock>>, in_flight: Arc, } +/// One Relay process and one start of that process. +#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +pub struct RelaySource { + pub runtime_id: u64, + pub incarnation: u64, +} + +#[derive(Debug, Default)] +struct SourceReport { + load: Option, + kv: Option, +} + #[derive(Debug)] struct Report { report: LoadReport, + received: Instant, expires: Instant, /// This process's requests in flight to the pool when the report arrived. ours: u64, @@ -77,42 +90,52 @@ impl PoolSignals { /// requests in. pub fn new(in_flight: Arc) -> Self { Self { - loads: DashMap::new(), - kv: DashMap::new(), + reports: DashMap::new(), regions: RwLock::default(), in_flight, } } /// Records what a pool's Relay reported. - pub fn record_load(&self, pool: &PoolId, report: LoadReport) { + pub fn record_load(&self, pool: &PoolId, source: RelaySource, report: LoadReport) { + let received = Instant::now(); let current = Report { - expires: Instant::now() + report.window * FRESH_WINDOWS, + expires: received + report.window * FRESH_WINDOWS, + received, report, ours: self.in_flight.get(pool), }; - self.loads.insert(pool.clone(), current); + let mut sources = self.reports.entry(pool.clone()).or_default(); + prune(&mut sources, received); + let source_report = sources.entry(source).or_default(); + source_report.load = Some(current); } /// Records the KV blocks a pool's workers hold for running requests, and /// their total, as the pool's Relay reported once a `window`. - pub fn record_kv(&self, pool: &PoolId, used: u64, total: u64, window: Duration) { - let expires = Instant::now() + window * FRESH_WINDOWS; - self.kv.insert( - pool.clone(), - KvReport { - used, - total, - expires, - }, - ); + pub fn record_kv( + &self, + pool: &PoolId, + source: RelaySource, + used: u64, + total: u64, + window: Duration, + ) { + let now = Instant::now(); + let mut sources = self.reports.entry(pool.clone()).or_default(); + prune(&mut sources, now); + let source_report = sources.entry(source).or_default(); + source_report.kv = Some(KvReport { + used, + total, + expires: now + window * FRESH_WINDOWS, + }); } /// Drops what a pool's Relay reported, for when the pool leaves or its /// Relay changes. pub fn forget(&self, pool: &PoolId) { - self.loads.remove(pool); - self.kv.remove(pool); + self.reports.remove(pool); } /// How full the pool is against its own capacity. `None` without a fresh @@ -127,10 +150,20 @@ impl PoolSignals { /// have no KV view yet; its slots alone then say how full it is. pub fn occupancy(&self, pool: &PoolId) -> Option { let now = Instant::now(); - let current = self - .loads - .get(pool) - .filter(|current| now <= current.expires)?; + let sources = self.reports.get(pool)?; + let (source, current) = sources + .iter() + .filter_map(|(source, source_report)| { + source_report + .load + .as_ref() + .filter(|current| { + now <= current.expires + && current.report.max_concurrency.is_some_and(|n| n > 0) + }) + .map(|current| (source, current)) + }) + .max_by_key(|(source, current)| (current.received, **source))?; let capacity = current .report .max_concurrency @@ -139,9 +172,9 @@ impl PoolSignals { .saturating_sub(current.ours) as f64; let generating = current.report.generating as f64; let slots = in_flight / capacity as f64; - let kv = self - .kv - .get(pool) + let kv = sources + .get(source) + .and_then(|source_report| source_report.kv.as_ref()) .filter(|kv| now <= kv.expires && kv.total > 0) .map(|kv| { let held = if generating > 0.0 { @@ -169,12 +202,24 @@ impl PoolSignals { } } +fn prune(sources: &mut HashMap, now: Instant) { + sources.retain(|_, report| { + report.load.as_ref().is_some_and(|load| now <= load.expires) + || report.kv.as_ref().is_some_and(|kv| now <= kv.expires) + }); +} + #[cfg(test)] mod tests { use crate::selector::tests::pool; use super::*; + const RELAY: RelaySource = RelaySource { + runtime_id: 1, + incarnation: 1, + }; + /// 64 requests in flight, 32 of them generating, of 128 slots, in a /// report fresh for a minute. fn half_full() -> LoadReport { @@ -191,7 +236,7 @@ mod tests { let in_flight = Arc::new(InFlightTracker::new()); let signals = PoolSignals::new(in_flight.clone()); let blue = pool("blue", "qwen", &["http://blue:8000"]).id; - signals.record_load(&blue, half_full()); + signals.record_load(&blue, RELAY, half_full()); // No KV view yet: the slots alone. let occupancy = signals.occupancy(&blue).unwrap(); assert_eq!((occupancy.load, occupancy.kv), (0.5, None)); @@ -199,7 +244,7 @@ mod tests { // Slot-bound: half the slots, and 1,000 of 16,384 KV blocks for 32 // running requests (2,000 for all 64). let window = Duration::from_secs(20); - signals.record_kv(&blue, 1000, 16384, window); + signals.record_kv(&blue, RELAY, 1000, 16384, window); let occupancy = signals.occupancy(&blue).unwrap(); assert_eq!(occupancy.slots, 0.5); assert!((occupancy.kv.unwrap() - 2000.0 / 16384.0).abs() < 1e-12); @@ -207,13 +252,14 @@ mod tests { // KV-bound: 8,000 blocks for 32 running requests; all 64 would need // 16,000, nearly every block, though half the slots are free. - signals.record_kv(&blue, 8000, 16384, window); + signals.record_kv(&blue, RELAY, 8000, 16384, window); let occupancy = signals.occupancy(&blue).unwrap(); assert!((occupancy.load - 16000.0 / 16384.0).abs() < 1e-12); // Without the workers' capacity, nothing is known. signals.record_load( &blue, + RELAY, LoadReport { max_concurrency: None, ..half_full() @@ -230,7 +276,7 @@ mod tests { let _before = in_flight.start(&blue); assert_eq!(signals.occupancy(&blue), None); - signals.record_load(&blue, half_full()); + signals.record_load(&blue, RELAY, half_full()); let _after = [in_flight.start(&blue), in_flight.start(&blue)]; assert!((signals.occupancy(&blue).unwrap().slots - 66.0 / 128.0).abs() < 1e-12); @@ -238,6 +284,7 @@ mod tests { let window = Duration::from_millis(50); signals.record_load( &blue, + RELAY, LoadReport { window, ..half_full() @@ -248,6 +295,83 @@ mod tests { assert_eq!(signals.occupancy(&blue), None); } + #[test] + fn replica_load_is_not_summed_and_kv_comes_from_the_selected_replica() { + let signals = PoolSignals::new(Arc::new(InFlightTracker::new())); + let blue = pool("blue", "qwen", &["http://blue:8000"]).id; + let green = pool("green", "qwen", &["http://green:8000"]).id; + let other = RelaySource { + runtime_id: 2, + incarnation: 1, + }; + let window = Duration::from_secs(20); + signals.record_load(&blue, RELAY, half_full()); + signals.record_kv(&blue, RELAY, 8000, 16384, window); + signals.record_load(&green, RELAY, half_full()); + signals.record_load( + &blue, + other, + LoadReport { + in_flight: 16, + generating: 8, + ..half_full() + }, + ); + signals.record_kv(&blue, other, 1000, 16384, window); + + let occupancy = signals.occupancy(&blue).unwrap(); + assert_eq!(occupancy.slots, 16.0 / 128.0); + assert_eq!(occupancy.kv, Some(2000.0 / 16384.0)); + assert_eq!(signals.occupancy(&green).unwrap().slots, 0.5); + + signals.record_load(&blue, other, half_full()); + assert_eq!(signals.occupancy(&blue).unwrap().slots, 0.5); + + signals.record_load( + &blue, + other, + LoadReport { + max_concurrency: None, + ..half_full() + }, + ); + assert_eq!(signals.occupancy(&blue).unwrap().slots, 0.5); + assert_eq!( + signals.occupancy(&blue).unwrap().kv, + Some(16000.0 / 16384.0) + ); + + let restarted = RelaySource { + runtime_id: RELAY.runtime_id, + incarnation: RELAY.incarnation + 1, + }; + signals.record_load(&blue, restarted, half_full()); + assert_eq!(signals.occupancy(&blue).unwrap().kv, None); + } + + #[test] + fn uses_a_surviving_replica_after_the_latest_report_expires() { + let signals = PoolSignals::new(Arc::new(InFlightTracker::new())); + let blue = pool("blue", "qwen", &["http://blue:8000"]).id; + let other = RelaySource { + runtime_id: 2, + incarnation: 1, + }; + signals.record_load(&blue, RELAY, half_full()); + signals.record_load( + &blue, + other, + LoadReport { + in_flight: 16, + window: Duration::from_millis(20), + ..half_full() + }, + ); + assert_eq!(signals.occupancy(&blue).unwrap().slots, 16.0 / 128.0); + std::thread::sleep(Duration::from_millis(100)); + assert_eq!(signals.occupancy(&blue).unwrap().slots, 0.5); + } + #[test] fn requests_that_finish_between_reports_stop_counting() { let in_flight = Arc::new(InFlightTracker::new()); @@ -255,7 +379,7 @@ mod tests { let blue = pool("blue", "qwen", &["http://blue:8000"]).id; // One of the report's 64 is this process's. let counted = in_flight.start(&blue); - signals.record_load(&blue, half_full()); + signals.record_load(&blue, RELAY, half_full()); // Short requests that start and finish before the next report. for _ in 0..100 {