From 518aefe11f9255f71fbdf41032a027b72df531ea Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Tue, 20 Oct 2026 17:20:11 -0700 Subject: [PATCH 1/5] feat(kv-router): expose session affinity replication for external hosts Signed-off-by: Jeffrey Wang --- lib/bindings/python/src/dynamo/_core.pyi | 7 +- lib/bindings/python/tests/test_kv_bindings.py | 165 ++++++- .../src/services/common/replica_sync.rs | 2 +- .../src/services/selection/affinity.rs | 71 ++- .../src/services/selection/core/operation.rs | 3 + .../services/selection/core/reservations.rs | 56 ++- .../src/services/selection/core/run.rs | 18 +- .../src/services/selection/core/tests.rs | 425 ++++++++++++++++-- .../src/services/selection/server.rs | 3 +- .../src/services/selection/service.rs | 4 +- lib/kv-router/src/services/selection/types.rs | 12 +- 11 files changed, 717 insertions(+), 49 deletions(-) diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index d7581bae659e..b428b69a0227 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -813,7 +813,7 @@ class SelectionService: """Select the best worker and book its load.""" ... - async def create_reservation(self, request: JsonLike) -> JsonLike: + async def create_reservation(self, request: JsonLike) -> Optional[JsonLike]: """Book a request's load against a worker, keyed by ``selection_id``. Without a ``worker_id``, replays the matching ``select``'s cached @@ -821,6 +821,11 @@ class SelectionService: other request fields are ignored. With a ``worker_id`` and the prompt, books explicitly under ``selection_id`` on that worker and discards any cached selection for the id. ``selection_id`` is required. + + Optional ``session_binding`` updates session affinity even if booking fails. + Set ``skip_load_booking=True`` for out-of-order delivery (completion + before reservation) or requests already tracked locally. This applies + only session affinity and returns ``None``. """ ... diff --git a/lib/bindings/python/tests/test_kv_bindings.py b/lib/bindings/python/tests/test_kv_bindings.py index e7643d8baf61..9931d54e15ad 100644 --- a/lib/bindings/python/tests/test_kv_bindings.py +++ b/lib/bindings/python/tests/test_kv_bindings.py @@ -291,20 +291,181 @@ async def test_selection_service_selects_registered_worker(monkeypatch): await service.shutdown_async() +@pytest.fixture +async def session_binding_services(monkeypatch): + monkeypatch.setenv("DYN_ROUTER_USE_KV_EVENTS", "false") + source = SelectionService(indexer_threads=1, session_affinity_ttl_secs=10) + peer = SelectionService(indexer_threads=1, session_affinity_ttl_secs=10) + try: + # Give both services the same routing partition and worker catalog. + for service in (source, peer): + for worker_id in (1, 2): + await service.upsert_worker( + { + "worker_id": worker_id, + "model_name": "model", + "routing_group": "host", + "endpoint": f"http://worker-{worker_id}", + "block_size": 4, + "max_num_batched_tokens": 1024, + } + ) + yield source, peer + finally: + await source.shutdown_async() + await peer.shutdown_async() + + @pytest.mark.skipif( SelectionService is None, reason="SelectionService requires the select-service Cargo feature", ) +@pytest.mark.parametrize("competing_score", ["load", "overlap"]) @pytest.mark.timeout(5) @pytest.mark.asyncio -async def test_selection_service_malformed_payload_raises_value_error(monkeypatch): +async def test_selection_service_affinity_scores( + session_binding_services, competing_score +): + source, peer = session_binding_services + partition = {"model_name": "model", "routing_group": "host"} + # Bind the session to worker 1 and replicate its assignment to the peer. + origin = await source.select_and_reserve( + { + **partition, + "token_ids": [1, 2, 3, 4], + "selection_id": "origin", + "session_id": "s", + "allowed_worker_ids": [1], + } + ) + await source.free_reservation("origin") + origin = json.loads(json.dumps(origin)) + await peer.create_reservation({**origin, "skip_load_booking": True}) + + # Make worker 2 preferable by load or cached-prefix overlap. + warm_tokens = list(range(100, 132)) + await peer.select_and_reserve( + { + **partition, + "token_ids": warm_tokens, + "selection_id": "background", + "allowed_worker_ids": [1 if competing_score == "load" else 2], + } + ) + if competing_score == "overlap": + await peer.free_reservation("background") + await peer.indexer_snapshot() + + # Unpinned requests prefer worker 2; session affinity keeps worker 1. + tokens = [5, 6, 7, 8] if competing_score == "load" else warm_tokens + request = {**partition, "token_ids": tokens} + control = await peer.select(request) + assert control["worker_id"] == 2 + selected = await peer.select_and_reserve( + {**request, "selection_id": "next-turn", "session_id": "s"} + ) + assert selected["worker_id"] == 1 + assert selected["session_binding"] == origin["session_binding"] + if competing_score == "overlap": + assert control["overlap"]["longest_matched"] == len(tokens) + assert selected["overlap"]["longest_matched"] == 0 + + +@pytest.mark.skipif( + SelectionService is None, + reason="SelectionService requires the select-service Cargo feature", +) +@pytest.mark.parametrize("booking", ["book", "skip", "fail"]) +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_selection_service_peer_affinity(session_binding_services, booking): + source, peer = session_binding_services + partition = {"model_name": "model", "routing_group": "host"} + # Export the session assignment chosen by the source. + origin = await source.select_and_reserve( + { + **partition, + "token_ids": [1, 2, 3, 4], + "selection_id": "origin", + "session_id": "s", + "allowed_worker_ids": [1], + } + ) + + # Import affinity with successful, skipped, or failed load booking. + request = {**origin, "skip_load_booking": booking == "skip"} + if booking == "fail": + request["dp_rank"] = 99 + with pytest.raises(SelectionServiceError): + await peer.create_reservation(request) + else: + result = await peer.create_reservation(request) + if booking == "skip": + assert result is None + else: + assert result["worker_id"] == 1 + + # Load reflects the booking outcome; the next turn keeps session affinity. + loads = peer.loads(**partition)[0]["loads"] + assert sum(row["active_requests"] for row in loads) == int(booking == "book") + selected = await peer.select_and_reserve( + { + **partition, + "selection_id": "next-turn", + "token_ids": [5, 6, 7, 8], + "session_id": "s", + } + ) + assert selected["worker_id"] == 1 + assert selected["session_binding"] == origin["session_binding"] + + +@pytest.mark.skipif( + SelectionService is None, + reason="SelectionService requires the select-service Cargo feature", +) +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_selection_service_skip_after_shutdown( + session_binding_services, +): + _, peer = session_binding_services + await peer.shutdown_async() + + # Skipping load booking must still enforce shutdown. + with pytest.raises(SelectionServiceError): + await peer.create_reservation( + {"selection_id": "after-shutdown", "skip_load_booking": True} + ) + + +@pytest.mark.skipif( + SelectionService is None, + reason="SelectionService requires the select-service Cargo feature", +) +@pytest.mark.parametrize( + "method,payload", + [ + pytest.param("upsert_worker", {"model_name": "model"}, id="worker"), + pytest.param( + "create_reservation", + {"selection_id": "malformed", "session_binding": {"session_id": "s"}}, + id="session-binding", + ), + ], +) +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_selection_service_malformed_payload_raises_value_error( + monkeypatch, method, payload +): monkeypatch.setenv("DYN_ROUTER_USE_KV_EVENTS", "false") service = SelectionService(indexer_threads=1) try: # Malformed input surfaces as a plain ValueError, not a SelectionServiceError. with pytest.raises(ValueError): - await service.upsert_worker({"model_name": "model"}) + await getattr(service, method)(payload) finally: service.shutdown() diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index 0fef46e5b038..292330154681 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -37,7 +37,7 @@ const AFFINITY_EVENT_CHANNEL_CAPACITY: usize = 4_096; /// One replicated session binding. #[derive(Debug, Clone, Serialize, Deserialize)] -pub(crate) struct AffinityBindingEvent { +pub struct AffinityBindingEvent { #[serde(flatten)] pub partition: RoutingPartitionId, pub session_id: String, diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index eab5b3578cbe..9dd0b9cb3c15 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -23,6 +23,9 @@ use tokio::sync::futures::OwnedNotified; use tokio::time::Instant; use tokio_util::sync::CancellationToken; +use crate::identity::RoutingPartitionId; +use crate::services::common::replica_sync::AffinityBindingEvent; + pub use crate::protocols::WorkerAffinityTarget as AffinityTarget; pub const MAX_SESSION_AFFINITY_TTL_SECS: u64 = 31_536_000; @@ -278,7 +281,7 @@ impl SessionAffinity { .unwrap_or_default() .as_nanos() as u64, ), - writer_id: AtomicU64::new(0), + writer_id: AtomicU64::new(fastrand::u64(..)), cancel: CancellationToken::new(), replica: OnceLock::new(), #[cfg(any(test, feature = "testing"))] @@ -480,9 +483,11 @@ impl SessionAffinity { lease: AffinityLease { coordinator: Arc::downgrade(&self.inner), session_id: session_id.to_string(), + binding_target: *target, revision: *revision, version: *version, active: true, + publish_binding_on_release: true, }, })) } @@ -534,9 +539,11 @@ impl SessionAffinity { let lease = AffinityLease { coordinator: Arc::downgrade(&self.inner), session_id: session_id.to_string(), + binding_target: target, revision, version, active: true, + publish_binding_on_release: true, }; lease.publish(target); Some(lease) @@ -605,6 +612,45 @@ impl SessionAffinity { Ok(Some(*target)) } + /// Keep session affinity alive while a mirrored reservation is active. + /// Only the matching, unexpired target/version can be leased; out-of-order + /// reservations must not keep a newer binding alive. Release refreshes the + /// local idle TTL without republishing the binding. + pub(super) fn acquire_mirrored_reservation_lease( + &self, + session_id: String, + expected_target: AffinityTarget, + expected_version: AffinityVersion, + ) -> Option { + let mut entry = self.inner.entries.get_mut(&session_id)?; + let AffinityEntry::Bound { + target, + version, + revision, + active_leases, + idle_deadline, + } = entry.value_mut() + else { + return None; + }; + if *target != expected_target + || *version != expected_version + || (*active_leases == 0 && *idle_deadline <= self.inner.now()) + { + return None; + } + *active_leases += 1; + Some(AffinityLease { + coordinator: Arc::downgrade(&self.inner), + session_id, + binding_target: *target, + revision: *revision, + version: *version, + active: true, + publish_binding_on_release: false, + }) + } + pub fn observe_replica_sequence(&self, sequence: u64) { self.inner.observe_replica_sequence(sequence); } @@ -846,9 +892,11 @@ impl AffinityInitialization { lease: AffinityLease { coordinator: Arc::downgrade(&inner), session_id: self.session_id.clone(), + binding_target: *bound, revision: self.revision, version: *version, active: true, + publish_binding_on_release: true, }, }; drop(entry); @@ -879,9 +927,11 @@ impl AffinityInitialization { let lease = AffinityLease { coordinator: Arc::downgrade(&inner), session_id: self.session_id.clone(), + binding_target: target, revision: self.revision, version, active: true, + publish_binding_on_release: true, }; lease.publish(target); Ok(lease) @@ -925,12 +975,26 @@ impl Drop for AffinityInitialization { pub struct AffinityLease { coordinator: Weak, session_id: String, + binding_target: AffinityTarget, revision: u64, version: AffinityVersion, active: bool, + publish_binding_on_release: bool, } impl AffinityLease { + /// Snapshot this booking's committed assignment, even if a peer changes the table. + pub(super) fn binding(&self, partition: &RoutingPartitionId) -> AffinityBindingEvent { + AffinityBindingEvent { + partition: partition.clone(), + session_id: self.session_id.clone(), + worker_id: self.binding_target.worker_id, + dp_rank: self.binding_target.dp_rank, + sequence: self.version.sequence, + writer_id: self.version.writer_id, + } + } + pub fn session_id(&self) -> &str { &self.session_id } @@ -965,6 +1029,7 @@ impl AffinityLease { let next_version = inner.next_version(); *current = target; *version = next_version; + self.binding_target = target; self.version = next_version; true } @@ -1013,7 +1078,9 @@ impl AffinityLease { *idle_deadline = inner.now() + inner.ttl; (*target, *version) }; - inner.publish_replica_update(&self.session_id, target, version); + if self.publish_binding_on_release { + inner.publish_replica_update(&self.session_id, target, version); + } } /// Drop the binding this lease holds (the bound worker is gone). A binding diff --git a/lib/kv-router/src/services/selection/core/operation.rs b/lib/kv-router/src/services/selection/core/operation.rs index 2817ee94f1ef..fff971a6712c 100644 --- a/lib/kv-router/src/services/selection/core/operation.rs +++ b/lib/kv-router/src/services/selection/core/operation.rs @@ -21,6 +21,7 @@ use crate::scheduling::{AdvisoryWorkerLoad, QueueRejection, SchedulingResponse, use super::super::error::SelectionError; use super::super::input::PromptView; +use crate::services::common::replica_sync::AffinityBindingEvent; /// Every input to one selection. Fields are independent: `session_context` /// feeds worker selection, `session` says what the core does with the session @@ -138,6 +139,8 @@ pub struct Selected { pub sequence_hashes: Option>, pub track_prefill_tokens: bool, pub effective_prefill_tokens: usize, + /// The versioned session assignment committed by this booking, if any. + pub session_binding: Option, pub kv_hint: Option, pub routing_hashes: Option>, pub shared_cache_hits: Option, diff --git a/lib/kv-router/src/services/selection/core/reservations.rs b/lib/kv-router/src/services/selection/core/reservations.rs index 4914d197381d..8e7b19acf8b5 100644 --- a/lib/kv-router/src/services/selection/core/reservations.rs +++ b/lib/kv-router/src/services/selection/core/reservations.rs @@ -25,6 +25,7 @@ struct ReservationBooking { /// Public block hashes to record into an approximate indexer once booked. routing_hashes: Option>, session_id: Option, + affinity_lease: Option, } /// `selection_id` -> the booking it holds. @@ -171,21 +172,52 @@ fn forget_reservation_if( impl SelectionCore { pub async fn create_reservation( &self, - req: ReservationRequest, - ) -> Result { + mut req: ReservationRequest, + ) -> Result, SelectionError> { self.ensure_running()?; + // Apply the peer's session affinity even if load booking is skipped or fails. + // Successful bookings retain a lease to keep affinity alive until completion. + let affinity_lease = if let Some(binding) = req.session_binding.take() { + let key = RoutingPartitionId::new(&req.model_name, &req.routing_group); + if binding.partition != key { + return Err(SelectionError::BadRequest( + "session binding must belong to the reservation's routing partition" + .to_string(), + )); + } + let session_id = binding.session_id.clone(); + let target = binding.target(); + let version = binding.version(); + self.dispatch_affinity_event(binding); + if req.skip_load_booking { + None + } else { + self.entry(&key).and_then(|entry| { + entry + .affinity + .get()? + .acquire_mirrored_reservation_lease(session_id, target, version) + }) + } + } else { + None + }; + if req.skip_load_booking { + return Ok(None); + } tokio::select! { biased; _ = self.cancel_token.cancelled() => { Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)) } - result = self.create_reservation_inner(req) => result, + result = self.create_reservation_inner(req, affinity_lease) => result.map(Some), } } async fn create_reservation_inner( &self, req: ReservationRequest, + affinity_lease: Option, ) -> Result { let key = RoutingPartitionId::new(req.model_name.clone(), req.routing_group.clone()); @@ -193,7 +225,9 @@ impl SelectionCore { // any cached selection for the id so a later replay can't book stale state. if let Some(worker_id) = req.worker_id { self.selection_cache.discard(&key, &req.selection_id); - return self.reserve_explicit(key, worker_id, req).await; + return self + .reserve_explicit(key, worker_id, req, affinity_lease) + .await; } // Replay form: peek, book, and consume only once the booking lands. A @@ -209,7 +243,10 @@ impl SelectionCore { req.selection_id ))); }; - let response = match self.book_cached_selection(pending, &req).await { + let response = match self + .book_cached_selection(pending, &req, affinity_lease) + .await + { Ok(response) => response, // The session moved: the cached worker can never be booked for it. Err(SelectionError::BadRequest(message)) => { @@ -230,6 +267,7 @@ impl SelectionCore { &self, pending: Arc, req: &ReservationRequest, + affinity_lease: Option, ) -> Result { let (entry, endpoint, prefill_load_hint) = self.resolve_cached_booking(&pending)?; let track_prefill_tokens = pending.track_prefill_tokens; @@ -247,6 +285,7 @@ impl SelectionCore { lora_name: pending.lora_name.clone(), routing_hashes: pending.routing_hashes.clone(), session_id: pending.session_id.clone(), + affinity_lease, }, ) .await @@ -299,6 +338,7 @@ impl SelectionCore { key: RoutingPartitionId, worker_id: WorkerId, req: ReservationRequest, + affinity_lease: Option, ) -> Result { let entry = self.ready_entry(&key)?; let normalized = req.prompt.view().normalize_for_reservation( @@ -355,6 +395,7 @@ impl SelectionCore { lora_name: req.prompt.lora_name, routing_hashes, session_id: None, + affinity_lease, }, ) .await @@ -380,6 +421,7 @@ impl SelectionCore { lora_name, routing_hashes, session_id, + affinity_lease, } = booking; let claim = self.claim_reservation(&selection_id, &key)?; @@ -387,7 +429,7 @@ impl SelectionCore { // initializer that may itself be waiting for this worker's capacity. let session = match (session_id.as_deref(), entry.affinity.get()) { (Some(session_id), Some(table)) => self - .hold_session(table, session_id, &key) + .hold_session(table, session_id, &key, None) .await? .map(|hold| (table, session_id, hold)), _ => None, @@ -410,7 +452,7 @@ impl SelectionCore { Some((table, session_id, hold)) => { self.commit_session(table, hold, session_id, worker, &key)? } - None => None, + None => affinity_lease, }; if let Some(hashes) = routing_hashes { self.record_routing_decision(&entry, worker, hashes).await; diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 27aceaa7cee9..9937c63dbd27 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -42,6 +42,8 @@ //! (`dropped_selection_future_frees_its_booking`, //! `dropped_book_selection_during_routing_record_frees_booking_and_claim`). +use std::collections::HashSet; + use super::super::affinity::{SessionAffinityMode, validate_dispatch_target}; use super::hint::{hint_capable_partition, transfer_hint_for_selection}; @@ -193,6 +195,7 @@ impl SelectionCore { sequence_hashes, track_prefill_tokens, effective_prefill_tokens, + session_binding, kv_hint, routing_hashes: _, shared_cache_hits: _, @@ -217,6 +220,7 @@ impl SelectionCore { .map(|threshold| load.prefill_load_exceeds(threshold)), }); SelectResponse { + session_binding, selection_id, sequence_hashes: sequence_hashes .map(|hashes| hashes.into_iter().map(|hash| hash as i64).collect()), @@ -306,7 +310,9 @@ impl SelectionCore { }; let affinity_target = match (&session, table) { (SessionBinding::Managed { session_id }, Some(table)) => { - affinity_hold = self.hold_session(table, session_id, &key).await?; + affinity_hold = self + .hold_session(table, session_id, &key, allowed_worker_ids.as_ref()) + .await?; affinity_hold.as_ref().and_then(Hold::target) } (SessionBinding::Query { session_id }, Some(table)) => table @@ -503,6 +509,7 @@ impl SelectionCore { // The routing hashes go to exactly one of: the reservation recorded // now, or the replay cache a later reservation records from. let mut routing_hashes = routing_hashes; + let mut session_binding = None; let booking = if let Some(claim) = claim { let Some(booking) = booking else { return Err(SelectionError::Internal( @@ -517,6 +524,7 @@ impl SelectionCore { } _ => None, }; + session_binding = affinity_lease.as_ref().map(|lease| lease.binding(&key)); if let Some(hashes) = routing_hashes.take() { self.record_routing_decision(&entry, response.best_worker, hashes) .await; @@ -558,6 +566,7 @@ impl SelectionCore { sequence_hashes: booked_sequence_hashes, track_prefill_tokens, effective_prefill_tokens: effective_prefill, + session_binding, kv_hint, routing_hashes: returned_routing_hashes, shared_cache_hits: host_shared_cache_hits, @@ -566,19 +575,22 @@ impl SelectionCore { } /// Hold `session_id` for a booking, re-initializing a binding whose target - /// this partition can no longer schedule. `None` when the table is full: a + /// this request can no longer schedule. `None` when the table is full: a /// router-side limit, not a client fault, so the request routes unpinned. pub(super) async fn hold_session( &self, table: &SessionAffinity, session_id: &str, key: &RoutingPartitionId, + allowed_worker_ids: Option<&HashSet>, ) -> Result, SelectionError> { let mut invalidations = 0; loop { match table.acquire(session_id, None).await { Ok(Hold::Bound { target, mut lease }) - if !self.catalog.is_schedulable(target, key) => + if !self.catalog.is_schedulable(target, key) + || allowed_worker_ids + .is_some_and(|ids| !ids.contains(&target.worker_id)) => { tracing::debug!( session_id, diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index d0f52b54ecc4..49d89d35fe5c 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -111,6 +111,16 @@ fn replay_reservation(selection_id: &str) -> ReservationRequest { expected_output_tokens: None, effective_prefill_tokens: None, track_prefill_tokens: None, + session_binding: None, + skip_load_booking: false, + } +} + +fn peer_affinity_update(binding: AffinityBindingEvent) -> ReservationRequest { + ReservationRequest { + session_binding: Some(binding), + skip_load_booking: true, + ..replay_reservation("peer-update") } } @@ -2880,7 +2890,7 @@ async fn joined_failover_preserves_the_successful_bookings_binding(#[case] commi let mut context = std::task::Context::from_waker(std::task::Waker::noop()); let mut r2_request = session_reservation("r2", "s"); - r2_request.allowed_worker_ids = Some(HashSet::from([replacement])); + r2_request.allowed_worker_ids = Some(HashSet::from([first.worker_id, replacement])); let mut r2 = Box::pin(core.select_and_reserve(r2_request)); assert!(r2.as_mut().poll(&mut context).is_pending()); core.catalog @@ -2925,11 +2935,19 @@ async fn late_failover_mismatch_preserves_the_committed_replacement() { let key = default_key(); let entry = core.entry(&key).expect("entry"); let table = entry.affinity.get().expect("affinity table"); - let old_hold = core.hold_session(table, "s", &key).await.unwrap().unwrap(); + let old_hold = core + .hold_session(table, "s", &key, None) + .await + .unwrap() + .unwrap(); core.catalog .set_lifecycle(first.worker_id, WorkerLifecycle::Draining, Vec::new()); let replacement = if first.worker_id == 1 { 2 } else { 1 }; - let new_hold = core.hold_session(table, "s", &key).await.unwrap().unwrap(); + let new_hold = core + .hold_session(table, "s", &key, None) + .await + .unwrap() + .unwrap(); let lease = core .commit_session( table, @@ -3039,7 +3057,8 @@ async fn failed_replay_preserves_a_newer_cached_selection() { let replacement = core .create_reservation(replay_reservation("pending")) .await - .expect("new selection survives the old replay failure"); + .expect("new selection survives the old replay failure") + .expect("reservation booked"); assert_eq!(replacement.worker_id, 2); } @@ -3084,45 +3103,391 @@ async fn free_releases_reservation_and_affinity_without_actor_progress() { .expect("completed cleanup must not prevent ID reuse"); } +#[rstest::rstest] +#[case::hard(SessionAffinityMode::Hard)] +#[case::soft(SessionAffinityMode::Soft)] #[tokio::test] -async fn hard_mode_rejects_dispatch_away_from_a_live_binding() { - let (core, first) = bound_session(SessionAffinityMode::Hard).await; +async fn session_rebinds_outside_allowlist(#[case] mode: SessionAffinityMode) { + let (core, first) = bound_session(mode).await; let other = if first.worker_id == 1 { 2 } else { 1 }; - - // Steering cannot reach the bound worker, so selection lands elsewhere. + // Excluding the bound worker permits failover in either affinity mode. let mut request = session_reservation("r2", "s"); request.allowed_worker_ids = Some(HashSet::from([other])); - let err = core - .select_and_reserve(request) - .await - .expect_err("hard affinity rejects a dispatch away from the binding"); - assert!(matches!(err, SelectionError::BadRequest(_)), "{err:?}"); - let entry = core.entry(&default_key()).expect("entry"); + let second = core.select_and_reserve(request).await.expect("rebind"); + assert_eq!(second.worker_id, other); + assert_eq!(bound_worker(&core, "s"), Some(other)); + assert_eq!(lease_count(&core, "s"), Some(1)); +} + +#[tokio::test] +async fn hard_affinity_rejects_policy_mismatch() { + // Bind to an eligible worker different from the custom policy's choice. + let (policy, _) = capturing_policy_factory(); + let core = core_with( + test_config(false), + SelectionHost::default(), + Some(policy), + WorkerType::Aggregated, + Some(SessionAffinityConfig::new(Duration::from_secs(10))), + ); + for id in [1, 2] { + core.upsert_worker(worker(id)).await.unwrap(); + } + let preferred = core.select(select_request()).await.unwrap().worker_id; + let bound = if preferred == 1 { 2 } else { 1 }; + core.dispatch_affinity_event(AffinityBindingEvent { + partition: default_key(), + session_id: "s".to_string(), + worker_id: bound, + dp_rank: Some(0), + sequence: 1, + writer_id: 99, + }); + assert_eq!(bound_worker(&core, "s"), Some(bound)); + // Hard affinity rejects the dispatch and releases its booking and binding. + let result = core + .select_and_reserve(session_reservation("next-turn", "s")) + .await; + assert!(matches!(result, Err(SelectionError::BadRequest(_)))); + let entry = core.entry(&default_key()).unwrap(); wait_until("rejected booking release", || { - !entry.scheduler.has_request("r2") + !entry.scheduler.has_request("next-turn") }) .await; assert!(core.reservation_index.read().is_empty()); - assert_eq!(bound_worker(&core, "s"), None, "stale binding is dropped"); + assert_eq!(bound_worker(&core, "s"), None); - // The session's next request re-initializes. - let mut request = session_reservation("r3", "s"); - request.allowed_worker_ids = Some(HashSet::from([other])); - let third = core.select_and_reserve(request).await.expect("rebind"); - assert_eq!(third.worker_id, other); - assert_eq!(bound_worker(&core, "s"), Some(other)); + // The next turn can initialize a fresh binding after the rejection. + let next = core + .select_and_reserve(session_reservation("retry", "s")) + .await + .expect("rebind after rejection"); + assert_eq!(next.worker_id, preferred); + assert_eq!(bound_worker(&core, "s"), Some(preferred)); +} + +struct CountingAffinitySink(AtomicUsize); + +impl super::super::affinity::AffinityReplicaSink for CountingAffinitySink { + fn publish( + &self, + _: &str, + _: WorkerAffinityTarget, + _: super::super::affinity::AffinityVersion, + ) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +#[tokio::test(start_paused = true)] +async fn binding_only_update_preserves_ttl() { + // Import an idle binding with a sink that detects accidental republication. + let (source, first) = bound_session(SessionAffinityMode::Hard).await; + let peer = core_with_session_affinity(); + peer.upsert_worker(worker(first.worker_id)).await.unwrap(); + let key = default_key(); + let publications = Arc::new(CountingAffinitySink(AtomicUsize::new(0))); + let entry = peer.entry(&key).unwrap(); + assert!( + entry + .affinity + .get() + .unwrap() + .enable_replication(99, publications.clone()) + ); + let loads_before = serde_json::to_value(peer.loads(None, None)).unwrap(); + let event = first.session_binding.unwrap(); + assert_eq!(event.worker_id, first.worker_id); + // Skipping load booking must leave the peer's loads unchanged. + assert!( + peer.create_reservation(peer_affinity_update(event)) + .await + .unwrap() + .is_none() + ); + assert_eq!( + serde_json::to_value(peer.loads(None, None)).unwrap(), + loads_before + ); + // Both idle bindings expire at the configured TTL without publishing. + tokio::time::advance(Duration::from_secs(9)).await; + assert!(bound_worker(&source, "s").is_some()); + assert!(bound_worker(&peer, "s").is_some()); + tokio::time::advance(Duration::from_secs(2)).await; + assert!(bound_worker(&source, "s").is_none()); + assert!(bound_worker(&peer, "s").is_none()); + assert_eq!(publications.0.load(Ordering::Relaxed), 0); +} + +#[tokio::test(start_paused = true)] +async fn mirrored_lease_preserves_ttl() { + // Mirror the source reservation and retain a lease on both services. + let source = core_with_session_affinity(); + let peer = core_with_session_affinity(); + for core in [&source, &peer] { + core.upsert_worker(worker(1)).await.unwrap(); + } + let key = default_key(); + let publications = Arc::new(CountingAffinitySink(AtomicUsize::new(0))); + let entry = peer.entry(&key).unwrap(); + assert!( + entry + .affinity + .get() + .unwrap() + .enable_replication(99, publications.clone()) + ); + + let selected = source + .select_and_reserve(session_reservation("turn", "s")) + .await + .unwrap(); + let binding = selected.session_binding.unwrap(); + peer.create_reservation(ReservationRequest { + worker_id: Some(selected.worker_id), + dp_rank: Some(selected.dp_rank), + prompt: prompt(), + session_binding: Some(binding.clone()), + ..replay_reservation("turn") + }) + .await + .unwrap() + .expect("mirrored booking"); + for core in [&source, &peer] { + assert_eq!(lease_count(core, "s"), Some(1)); + } + let mirrored = peer + .reservation_index + .read() + .get("turn") + .unwrap() + ._affinity_lease + .as_ref() + .unwrap() + .binding(&key); + assert_eq!(mirrored.version(), binding.version()); + + // A running request keeps both bindings alive beyond the idle TTL. + tokio::time::advance(Duration::from_secs(11)).await; + for core in [&source, &peer] { + assert_eq!( + bound_worker(core, "s"), + Some(1), + "an active turn cannot expire" + ); + core.free_reservation("turn").await.unwrap(); + assert_eq!(lease_count(core, "s"), Some(0)); + } + // Completion starts a full idle TTL; the peer's lease must not republish. + tokio::time::advance(Duration::from_secs(9)).await; + for core in [&source, &peer] { + assert_eq!( + bound_worker(core, "s"), + Some(1), + "idle TTL starts at completion" + ); + } + tokio::time::advance(Duration::from_secs(2)).await; + for core in [&source, &peer] { + assert_eq!(bound_worker(core, "s"), None); + } + assert_eq!(publications.0.load(Ordering::Relaxed), 0); } #[tokio::test] -async fn soft_mode_follows_the_dispatch() { - let (core, first) = bound_session(SessionAffinityMode::Soft).await; - let other = if first.worker_id == 1 { 2 } else { 1 }; +async fn peer_bindings_converge() { + // Create competing assignments from independent services. + let source = core_with_session_affinity(); + let peer = core_with_session_affinity(); + for core in [&source, &peer] { + for id in [1, 2, 3] { + core.upsert_worker(worker(id)).await.unwrap(); + } + } + let mut bindings = Vec::new(); + for (core, worker_id) in [(&source, 1), (&peer, 2)] { + let mut request = session_reservation("first", "s"); + request.allowed_worker_ids = Some(HashSet::from([worker_id])); + let selected = core.select_and_reserve(request).await.unwrap(); + bindings.push(selected.session_binding.unwrap()); + } + let mut first = bindings.remove(0); + let mut concurrent = bindings.remove(0); + // Exercise a sequence tie between independent writers. + let sequence = first.sequence.max(concurrent.sequence); + first.sequence = sequence; + concurrent.sequence = sequence; + assert_ne!(first.writer_id, concurrent.writer_id); + let winner = if first.version() > concurrent.version() { + first.clone() + } else { + concurrent.clone() + }; + // Opposite delivery orders must choose the same version winner. + for (core, events) in [ + (&source, [&first, &concurrent]), + (&peer, [&concurrent, &first]), + ] { + for event in events { + core.create_reservation(peer_affinity_update(event.clone())) + .await + .unwrap(); + } + } + for core in [&source, &peer] { + assert_eq!(bound_worker(core, "s"), Some(winner.worker_id)); + } - let mut request = session_reservation("r2", "s"); - request.allowed_worker_ids = Some(HashSet::from([other])); - let second = core.select_and_reserve(request).await.expect("soft rebind"); - assert_eq!(second.worker_id, other); - assert_eq!(bound_worker(&core, "s"), Some(other)); + // A later rebind wins even when older updates are replayed afterward. + let mut request = session_reservation("rebind", "s"); + request.allowed_worker_ids = Some(HashSet::from([3])); + let selected = source.select_and_reserve(request).await.unwrap(); + let newer = selected.session_binding.unwrap(); + assert!(newer.version() > winner.version()); + peer.create_reservation(peer_affinity_update(newer)) + .await + .unwrap(); + for core in [&source, &peer] { + for event in [&first, &concurrent] { + core.create_reservation(peer_affinity_update(event.clone())) + .await + .unwrap(); + } + assert_eq!(bound_worker(core, "s"), Some(3)); + } +} + +#[rstest::rstest] +#[case::book(false, false)] +#[case::skip(true, false)] +#[case::fail(false, true)] +#[tokio::test] +async fn affinity_survives_booking_outcomes( + #[case] skip_load_booking: bool, + #[case] invalid_rank: bool, +) { + // Import the binding with successful, skipped, or failed load booking. + let (_, first) = bound_session(SessionAffinityMode::Hard).await; + let peer = core_with_session_affinity(); + for id in [1, 2] { + peer.upsert_worker(worker(id)).await.unwrap(); + } + let binding = first.session_binding.unwrap(); + let result = peer + .create_reservation(ReservationRequest { + worker_id: Some(first.worker_id), + dp_rank: Some(if invalid_rank { 99 } else { 0 }), + prompt: prompt(), + session_binding: Some(binding.clone()), + skip_load_booking, + ..replay_reservation("mirror") + }) + .await; + if invalid_rank { + assert!(result.is_err()); + } else { + assert_eq!(result.unwrap().is_none(), skip_load_booking); + } + // Every outcome imports affinity; only successful bookings retain load/leases. + assert_eq!(bound_worker(&peer, "s"), Some(first.worker_id)); + assert_eq!( + lease_count(&peer, "s"), + Some(usize::from(!skip_load_booking && !invalid_rank)) + ); + let active: usize = peer + .loads(None, None) + .iter() + .flat_map(|partition| &partition.loads) + .map(|load| load.active_requests) + .sum(); + assert_eq!(active, usize::from(!skip_load_booking && !invalid_rank)); + // The next turn reuses the imported assignment and version. + let next = peer + .select_and_reserve(session_reservation("next", "s")) + .await + .unwrap(); + let next_binding = next.session_binding.unwrap(); + assert_eq!(next_binding.worker_id, binding.worker_id); + assert_eq!(next_binding.version(), binding.version()); +} + +#[tokio::test] +async fn binding_rejects_wrong_partition() { + let (core, first) = bound_session(SessionAffinityMode::Hard).await; + // Attach another model's binding to this partition's reservation. + let mut binding = first.session_binding.unwrap(); + binding.partition = RoutingPartitionId::new("other", "default"); + let request = ReservationRequest { + worker_id: Some(first.worker_id), + dp_rank: Some(first.dp_rank), + prompt: prompt(), + session_binding: Some(binding), + ..replay_reservation("mismatch") + }; + // Reject the mismatch before creating a load booking. + assert!(matches!( + core.create_reservation(request).await, + Err(SelectionError::BadRequest(_)) + )); + assert!(core.reservation_index.read().is_empty()); + assert!( + !core + .entry(&default_key()) + .unwrap() + .scheduler + .has_request("mismatch") + ); +} + +#[rstest::rstest] +#[case::different_worker(false)] +#[case::same_worker(true)] +#[tokio::test] +async fn stale_binding_preserves_load(#[case] return_to_original_worker: bool) { + // Rebind an active session, optionally returning to its original worker. + let source = core_with_session_affinity(); + let peer = core_with_session_affinity(); + for core in [&source, &peer] { + for id in [1, 2] { + core.upsert_worker(worker(id)).await.unwrap(); + } + } + let mut old_request = session_reservation("old", "s"); + old_request.allowed_worker_ids = Some(HashSet::from([1])); + let old = source.select_and_reserve(old_request).await.unwrap(); + let mut new_request = session_reservation("new", "s"); + new_request.allowed_worker_ids = Some(HashSet::from([2])); + let mut new = source.select_and_reserve(new_request).await.unwrap(); + if return_to_original_worker { + let mut returning = session_reservation("returning", "s"); + returning.allowed_worker_ids = Some(HashSet::from([1])); + new = source.select_and_reserve(returning).await.unwrap(); + } + let replacement = new.worker_id; + // Deliver the newest binding before the old request's delayed reservation. + peer.create_reservation(peer_affinity_update(new.session_binding.unwrap())) + .await + .unwrap(); + peer.create_reservation(ReservationRequest { + worker_id: Some(old.worker_id), + dp_rank: Some(old.dp_rank), + prompt: prompt(), + session_binding: old.session_binding, + ..replay_reservation("old") + }) + .await + .unwrap() + .expect("old request still needs load accounting"); + // Count the original request without reverting or leasing the newer binding. + assert_eq!(bound_worker(&peer, "s"), Some(replacement)); + let loads = peer.loads(None, None); + let original_worker = loads + .iter() + .flat_map(|p| &p.loads) + .find(|w| w.worker_id == 1) + .unwrap(); + assert_eq!(original_worker.active_requests, 1); + assert_eq!(lease_count(&peer, "s"), Some(0)); } #[tokio::test] diff --git a/lib/kv-router/src/services/selection/server.rs b/lib/kv-router/src/services/selection/server.rs index cfb6af7ab39f..fa9e4a8d1360 100644 --- a/lib/kv-router/src/services/selection/server.rs +++ b/lib/kv-router/src/services/selection/server.rs @@ -129,7 +129,8 @@ async fn create_reservation( Err(error) => return json_rejection(error), }; match state.service.create_reservation(req).await { - Ok(response) => (StatusCode::CREATED, Json(response)).into_response(), + Ok(Some(response)) => (StatusCode::CREATED, Json(response)).into_response(), + Ok(None) => StatusCode::NO_CONTENT.into_response(), Err(error) => error.into_response(), } } diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 057b328dcbf5..0fb891028a98 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -406,10 +406,12 @@ impl SelectionService { .await } + /// Apply an optional peer binding, then book load unless `skip_load_booking` is set. + /// Affinity is retained even if booking fails; skipped bookings return `None`. pub async fn create_reservation( &self, req: ReservationRequest, - ) -> Result { + ) -> Result, SelectionError> { self.core.create_reservation(req).await } diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index 26a51a93e225..fcffa89f8884 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -15,6 +15,7 @@ use crate::protocols::{ use crate::scheduling::config::RouterConfigOverride; pub use crate::scheduling::{OverlapScoresResponse, SharedCacheOverlapScore, WorkerOverlapScore}; use crate::scheduling::{PotentialLoad, SessionContext, WorkerSelectionInputTrigger}; +use crate::services::common::replica_sync::AffinityBindingEvent; use crate::services::overlap::MooncakeOverlapSummary; use super::input::PromptRequest; @@ -551,7 +552,8 @@ fn resolve_session_context( /// Booking request: replay the selection cached under `selection_id`, or book /// self-contained with `worker_id`. The replay books exactly what `select` captured; -/// request fields other than the ids and model/routing-group are ignored. +/// booking fields other than the ids and model/routing-group are ignored on replay. +/// A peer's session binding is applied independently, including when booking is skipped. #[derive(Debug, Deserialize)] pub struct ReservationRequest { #[serde(default = "default_model_name")] @@ -571,6 +573,11 @@ pub struct ReservationRequest { pub expected_output_tokens: Option, pub effective_prefill_tokens: Option, pub track_prefill_tokens: Option, + /// Assignment received from a peer; applied before booking, without republishing. + pub session_binding: Option, + /// Skip load booking for requests already tracked or completed locally. + #[serde(default)] + pub skip_load_booking: bool, } #[derive(Debug, Deserialize)] @@ -602,6 +609,9 @@ pub struct OverlapScoresRequest { #[derive(Debug, Serialize)] pub struct SelectResponse { + /// Assignment to replicate alongside the reservation; absent for unpinned requests. + #[serde(skip_serializing_if = "Option::is_none")] + pub session_binding: Option, #[serde(skip_serializing_if = "Option::is_none")] pub selection_id: Option, #[serde(skip_serializing_if = "Option::is_none")] From f13adc11476263e40874f434ea9d707321b964f3 Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Sat, 24 Oct 2026 21:20:22 -0700 Subject: [PATCH 2/5] refactor(kv-router): stream host session bindings and simplify affinity replication Signed-off-by: Jeffrey Wang --- lib/bindings/python/rust/lib.rs | 2 + lib/bindings/python/rust/llm/kv.rs | 64 ++- lib/bindings/python/src/dynamo/_core.pyi | 33 +- lib/bindings/python/tests/test_kv_bindings.py | 144 ++--- .../src/services/common/replica_sync.rs | 10 +- .../src/services/selection/affinity.rs | 71 +-- .../src/services/selection/core/mod.rs | 50 +- .../src/services/selection/core/operation.rs | 3 - .../services/selection/core/reservations.rs | 54 +- .../src/services/selection/core/run.rs | 12 +- .../src/services/selection/core/tests.rs | 507 +++++------------- .../src/services/selection/core/workers.rs | 1 + lib/kv-router/src/services/selection/mod.rs | 7 +- .../src/services/selection/server.rs | 3 +- .../src/services/selection/service.rs | 108 +++- lib/kv-router/src/services/selection/tests.rs | 22 + lib/kv-router/src/services/selection/types.rs | 12 +- lib/llm/src/kv_router/embedded.rs | 1 + 18 files changed, 473 insertions(+), 631 deletions(-) diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index f0184a01f4e1..d451d9baf1e6 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -360,6 +360,8 @@ fn register_core(m: &Bound<'_, PyModule>) -> PyResult<()> { #[cfg(feature = "select-service")] m.add_class::()?; #[cfg(feature = "select-service")] + m.add_class::()?; + #[cfg(feature = "select-service")] m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 81946cabd7f6..f20e969c88cd 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -657,12 +657,36 @@ pub(crate) struct SelectionService { inner: Arc, } +#[cfg(feature = "select-service")] +#[pyclass] +pub(crate) struct SessionBindingStream { + updates: Arc, +} + +#[cfg(feature = "select-service")] +#[pymethods] +impl SessionBindingStream { + fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __anext__<'p>(&self, py: Python<'p>) -> PyResult> { + let updates = self.updates.clone(); + crate::future_into_py(py, async move { + let event = updates.recv().await.ok_or_else(|| { + pyo3::exceptions::PyStopAsyncIteration::new_err("selection service shut down") + })?; + Python::with_gil(|py| pythonize(py, &event).map(|o| o.unbind()).map_err(to_pyerr)) + }) + } +} + #[cfg(feature = "select-service")] #[pymethods] impl SelectionService { /// Create a selection service. `indexer_threads` sizes the KV indexer pool. #[new] - #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, session_affinity_ttl_secs = None))] + #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, session_affinity_ttl_secs = None, session_affinity_mode = "hard", session_replication_writer_id = None))] #[allow(clippy::too_many_arguments)] fn new( py: Python<'_>, @@ -672,7 +696,12 @@ impl SelectionService { replica_sync_peers: Option>, selection_cache: Option, session_affinity_ttl_secs: Option, + session_affinity_mode: &str, + session_replication_writer_id: Option, ) -> PyResult { + let session_affinity_mode = session_affinity_mode + .parse::() + .map_err(PyValueError::new_err)?; let replica_sync_peers = replica_sync_peers.unwrap_or_default(); if replica_sync_port.is_none() && !replica_sync_peers.is_empty() { return Err(PyValueError::new_err( @@ -690,7 +719,8 @@ impl SelectionService { ) .indexer_threads(indexer_threads) .indexer_peers(indexer_peers.unwrap_or_default()) - .selection_cache(selection_cache.unwrap_or_default().inner); + .selection_cache(selection_cache.unwrap_or_default().inner) + .session_affinity_mode(session_affinity_mode); if let Some(port) = replica_sync_port { builder = builder.replica_sync(port, replica_sync_peers); } @@ -699,6 +729,9 @@ impl SelectionService { session_affinity_ttl_from_secs(ttl).map_err(PyValueError::new_err)?, ); } + if let Some(writer_id) = session_replication_writer_id { + builder = builder.session_binding_replication(writer_id); + } let inner = py .allow_threads(|| crate::bridge_runtime().block_on(builder.build())) .map_err(to_pyerr)?; @@ -717,6 +750,33 @@ impl SelectionService { py.allow_threads(|| crate::bridge_runtime().block_on(service.shutdown())); } + /// Session bindings this replica commits or refreshes, for the host to + /// broadcast to its peers. Ends when the service shuts down. + fn session_binding_updates(&self) -> PyResult { + Ok(SessionBindingStream { + updates: Arc::new( + self.inner + .session_binding_updates() + .map_err(selection_to_pyerr)?, + ), + }) + } + + fn apply_session_binding<'p>( + &self, + py: Python<'p>, + update: PyObject, + ) -> PyResult> { + let event: selection::AffinityBindingEvent = + depythonize(update.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?; + let service = self.inner.clone(); + crate::future_into_py(py, async move { + service + .apply_session_binding(event) + .map_err(selection_to_pyerr) + }) + } + /// Await service shutdown without blocking the Python event loop. fn shutdown_async<'p>(&self, py: Python<'p>) -> PyResult> { let service = Arc::clone(&self.inner); diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index b428b69a0227..f13fb294a29f 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -765,11 +765,17 @@ class SelectionService: replica_sync_peers: Optional[list[str]] = None, selection_cache: Optional[SelectionCacheConfig] = None, session_affinity_ttl_secs: Optional[float] = None, + session_affinity_mode: Literal["hard", "soft"] = "hard", + session_replication_writer_id: Optional[int] = None, ) -> None: """Create a selection service. `indexer_threads` sizes the KV indexer pool. `session_affinity_ttl_secs` enables session affinity with an idle TTL between 1 and 31536000 seconds. + `session_affinity_mode`: hard rejects dispatch away from a live binding; + soft rebinds when the bound worker is excluded from the request's candidates. + `session_replication_writer_id` exports affinity updates for a host transport; + use an id unique to this replica incarnation. """ ... @@ -783,6 +789,26 @@ class SelectionService: """ ... + async def shutdown_async(self) -> None: + """Await service shutdown without blocking the Python event loop.""" + ... + + def session_binding_updates(self) -> AsyncIterator[JsonLike]: + """Stream session bindings this service commits or refreshes, for the + host to broadcast to its peers. + + Requires session_replication_writer_id. Ends on shutdown. + """ + ... + + async def apply_session_binding(self, update: JsonLike) -> None: + """Merge a versioned peer binding without echoing it. + + A binding this replica cannot place yet is dropped; the owner's next + refresh repairs it. + """ + ... + async def upsert_worker(self, worker: JsonLike) -> JsonLike: """Upsert a worker and subscribe to its live KV events; returns its catalog record.""" ... @@ -813,7 +839,7 @@ class SelectionService: """Select the best worker and book its load.""" ... - async def create_reservation(self, request: JsonLike) -> Optional[JsonLike]: + async def create_reservation(self, request: JsonLike) -> JsonLike: """Book a request's load against a worker, keyed by ``selection_id``. Without a ``worker_id``, replays the matching ``select``'s cached @@ -821,11 +847,6 @@ class SelectionService: other request fields are ignored. With a ``worker_id`` and the prompt, books explicitly under ``selection_id`` on that worker and discards any cached selection for the id. ``selection_id`` is required. - - Optional ``session_binding`` updates session affinity even if booking fails. - Set ``skip_load_booking=True`` for out-of-order delivery (completion - before reservation) or requests already tracked locally. This applies - only session affinity and returns ``None``. """ ... diff --git a/lib/bindings/python/tests/test_kv_bindings.py b/lib/bindings/python/tests/test_kv_bindings.py index 9931d54e15ad..ad014bf18ca5 100644 --- a/lib/bindings/python/tests/test_kv_bindings.py +++ b/lib/bindings/python/tests/test_kv_bindings.py @@ -294,10 +294,18 @@ async def test_selection_service_selects_registered_worker(monkeypatch): @pytest.fixture async def session_binding_services(monkeypatch): monkeypatch.setenv("DYN_ROUTER_USE_KV_EVENTS", "false") - source = SelectionService(indexer_threads=1, session_affinity_ttl_secs=10) - peer = SelectionService(indexer_threads=1, session_affinity_ttl_secs=10) + source = SelectionService( + indexer_threads=1, + session_affinity_ttl_secs=10, + session_replication_writer_id=11, + ) + peer = SelectionService( + indexer_threads=1, + session_affinity_ttl_secs=10, + session_affinity_mode="soft", + session_replication_writer_id=22, + ) try: - # Give both services the same routing partition and worker catalog. for service in (source, peer): for worker_id in (1, 2): await service.upsert_worker( @@ -320,16 +328,16 @@ async def session_binding_services(monkeypatch): SelectionService is None, reason="SelectionService requires the select-service Cargo feature", ) -@pytest.mark.parametrize("competing_score", ["load", "overlap"]) @pytest.mark.timeout(5) @pytest.mark.asyncio -async def test_selection_service_affinity_scores( - session_binding_services, competing_score -): +async def test_selection_service_session_binding_stream(session_binding_services): source, peer = session_binding_services + stream = source.session_binding_updates() + peer_stream = peer.session_binding_updates() + with pytest.raises(SelectionServiceError): + source.session_binding_updates() partition = {"model_name": "model", "routing_group": "host"} - # Bind the session to worker 1 and replicate its assignment to the peer. - origin = await source.select_and_reserve( + selected = await source.select_and_reserve( { **partition, "token_ids": [1, 2, 3, 4], @@ -338,105 +346,41 @@ async def test_selection_service_affinity_scores( "allowed_worker_ids": [1], } ) + assert "session_binding" not in selected + binding = json.loads(json.dumps(await anext(stream))) + assert binding["writer_id"] == 11 + await peer.apply_session_binding(binding) await source.free_reservation("origin") - origin = json.loads(json.dumps(origin)) - await peer.create_reservation({**origin, "skip_load_booking": True}) - - # Make worker 2 preferable by load or cached-prefix overlap. - warm_tokens = list(range(100, 132)) - await peer.select_and_reserve( - { - **partition, - "token_ids": warm_tokens, - "selection_id": "background", - "allowed_worker_ids": [1 if competing_score == "load" else 2], - } - ) - if competing_score == "overlap": - await peer.free_reservation("background") - await peer.indexer_snapshot() - - # Unpinned requests prefer worker 2; session affinity keeps worker 1. - tokens = [5, 6, 7, 8] if competing_score == "load" else warm_tokens - request = {**partition, "token_ids": tokens} - control = await peer.select(request) - assert control["worker_id"] == 2 - selected = await peer.select_and_reserve( - {**request, "selection_id": "next-turn", "session_id": "s"} - ) - assert selected["worker_id"] == 1 - assert selected["session_binding"] == origin["session_binding"] - if competing_score == "overlap": - assert control["overlap"]["longest_matched"] == len(tokens) - assert selected["overlap"]["longest_matched"] == 0 - - -@pytest.mark.skipif( - SelectionService is None, - reason="SelectionService requires the select-service Cargo feature", -) -@pytest.mark.parametrize("booking", ["book", "skip", "fail"]) -@pytest.mark.timeout(5) -@pytest.mark.asyncio -async def test_selection_service_peer_affinity(session_binding_services, booking): - source, peer = session_binding_services - partition = {"model_name": "model", "routing_group": "host"} - # Export the session assignment chosen by the source. - origin = await source.select_and_reserve( - { - **partition, - "token_ids": [1, 2, 3, 4], - "selection_id": "origin", - "session_id": "s", - "allowed_worker_ids": [1], - } - ) - - # Import affinity with successful, skipped, or failed load booking. - request = {**origin, "skip_load_booking": booking == "skip"} - if booking == "fail": - request["dp_rank"] = 99 - with pytest.raises(SelectionServiceError): - await peer.create_reservation(request) - else: - result = await peer.create_reservation(request) - if booking == "skip": - assert result is None - else: - assert result["worker_id"] == 1 - - # Load reflects the booking outcome; the next turn keeps session affinity. - loads = peer.loads(**partition)[0]["loads"] - assert sum(row["active_requests"] for row in loads) == int(booking == "book") + assert await anext(stream) == binding # A repeated version refreshes TTL. selected = await peer.select_and_reserve( { **partition, - "selection_id": "next-turn", "token_ids": [5, 6, 7, 8], + "selection_id": "next", "session_id": "s", } ) assert selected["worker_id"] == 1 - assert selected["session_binding"] == origin["session_binding"] - - -@pytest.mark.skipif( - SelectionService is None, - reason="SelectionService requires the select-service Cargo feature", -) -@pytest.mark.timeout(5) -@pytest.mark.asyncio -async def test_selection_service_skip_after_shutdown( - session_binding_services, -): - _, peer = session_binding_services - await peer.shutdown_async() - - # Skipping load booking must still enforce shutdown. + assert ( + await anext(peer_stream) == binding + ) # Local commits preserve imported versions. + # The default hard mode rejects an excluded binding; soft mode rebinds. + request = { + **partition, + "token_ids": [9, 10, 11, 12], + "selection_id": "excluded", + "session_id": "s", + "allowed_worker_ids": [2], + } + with pytest.raises(ValueError, match="session s is bound to worker 1, not 2"): + await source.select_and_reserve(request) + assert (await peer.select_and_reserve(request))["worker_id"] == 2 + assert (await anext(peer_stream))["worker_id"] == 2 + await source.shutdown_async() + with pytest.raises(StopAsyncIteration): + await anext(stream) with pytest.raises(SelectionServiceError): - await peer.create_reservation( - {"selection_id": "after-shutdown", "skip_load_booking": True} - ) + await source.apply_session_binding(binding) @pytest.mark.skipif( @@ -448,8 +392,8 @@ async def test_selection_service_skip_after_shutdown( [ pytest.param("upsert_worker", {"model_name": "model"}, id="worker"), pytest.param( - "create_reservation", - {"selection_id": "malformed", "session_binding": {"session_id": "s"}}, + "apply_session_binding", + {"session_id": "s"}, id="session-binding", ), ], diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index 292330154681..ffe819a252d3 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -33,7 +33,7 @@ const REPLICA_TOPIC: &[u8] = b"dynamo.slot-tracker.v1"; /// Session-affinity bindings ride the same mesh on their own topic, so peers /// that do not subscribe to it never see them. const AFFINITY_TOPIC: &[u8] = b"dynamo.session-affinity.v1"; -const AFFINITY_EVENT_CHANNEL_CAPACITY: usize = 4_096; +pub(crate) const AFFINITY_EVENT_CHANNEL_CAPACITY: usize = 4_096; /// One replicated session binding. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -63,10 +63,10 @@ impl AffinityBindingEvent { #[cfg(feature = "standalone-selection")] /// Publishes bindings from a [`super::super::selection::affinity::SessionAffinity`] -/// into the mesh. Best effort: a full channel drops the update. -struct AffinityMeshSink { - partition: RoutingPartitionId, - tx: mpsc::Sender, +/// into a transport channel. Best effort: a full channel drops the update. +pub(crate) struct AffinityMeshSink { + pub(crate) partition: RoutingPartitionId, + pub(crate) tx: mpsc::Sender, } #[cfg(feature = "standalone-selection")] diff --git a/lib/kv-router/src/services/selection/affinity.rs b/lib/kv-router/src/services/selection/affinity.rs index 9dd0b9cb3c15..eab5b3578cbe 100644 --- a/lib/kv-router/src/services/selection/affinity.rs +++ b/lib/kv-router/src/services/selection/affinity.rs @@ -23,9 +23,6 @@ use tokio::sync::futures::OwnedNotified; use tokio::time::Instant; use tokio_util::sync::CancellationToken; -use crate::identity::RoutingPartitionId; -use crate::services::common::replica_sync::AffinityBindingEvent; - pub use crate::protocols::WorkerAffinityTarget as AffinityTarget; pub const MAX_SESSION_AFFINITY_TTL_SECS: u64 = 31_536_000; @@ -281,7 +278,7 @@ impl SessionAffinity { .unwrap_or_default() .as_nanos() as u64, ), - writer_id: AtomicU64::new(fastrand::u64(..)), + writer_id: AtomicU64::new(0), cancel: CancellationToken::new(), replica: OnceLock::new(), #[cfg(any(test, feature = "testing"))] @@ -483,11 +480,9 @@ impl SessionAffinity { lease: AffinityLease { coordinator: Arc::downgrade(&self.inner), session_id: session_id.to_string(), - binding_target: *target, revision: *revision, version: *version, active: true, - publish_binding_on_release: true, }, })) } @@ -539,11 +534,9 @@ impl SessionAffinity { let lease = AffinityLease { coordinator: Arc::downgrade(&self.inner), session_id: session_id.to_string(), - binding_target: target, revision, version, active: true, - publish_binding_on_release: true, }; lease.publish(target); Some(lease) @@ -612,45 +605,6 @@ impl SessionAffinity { Ok(Some(*target)) } - /// Keep session affinity alive while a mirrored reservation is active. - /// Only the matching, unexpired target/version can be leased; out-of-order - /// reservations must not keep a newer binding alive. Release refreshes the - /// local idle TTL without republishing the binding. - pub(super) fn acquire_mirrored_reservation_lease( - &self, - session_id: String, - expected_target: AffinityTarget, - expected_version: AffinityVersion, - ) -> Option { - let mut entry = self.inner.entries.get_mut(&session_id)?; - let AffinityEntry::Bound { - target, - version, - revision, - active_leases, - idle_deadline, - } = entry.value_mut() - else { - return None; - }; - if *target != expected_target - || *version != expected_version - || (*active_leases == 0 && *idle_deadline <= self.inner.now()) - { - return None; - } - *active_leases += 1; - Some(AffinityLease { - coordinator: Arc::downgrade(&self.inner), - session_id, - binding_target: *target, - revision: *revision, - version: *version, - active: true, - publish_binding_on_release: false, - }) - } - pub fn observe_replica_sequence(&self, sequence: u64) { self.inner.observe_replica_sequence(sequence); } @@ -892,11 +846,9 @@ impl AffinityInitialization { lease: AffinityLease { coordinator: Arc::downgrade(&inner), session_id: self.session_id.clone(), - binding_target: *bound, revision: self.revision, version: *version, active: true, - publish_binding_on_release: true, }, }; drop(entry); @@ -927,11 +879,9 @@ impl AffinityInitialization { let lease = AffinityLease { coordinator: Arc::downgrade(&inner), session_id: self.session_id.clone(), - binding_target: target, revision: self.revision, version, active: true, - publish_binding_on_release: true, }; lease.publish(target); Ok(lease) @@ -975,26 +925,12 @@ impl Drop for AffinityInitialization { pub struct AffinityLease { coordinator: Weak, session_id: String, - binding_target: AffinityTarget, revision: u64, version: AffinityVersion, active: bool, - publish_binding_on_release: bool, } impl AffinityLease { - /// Snapshot this booking's committed assignment, even if a peer changes the table. - pub(super) fn binding(&self, partition: &RoutingPartitionId) -> AffinityBindingEvent { - AffinityBindingEvent { - partition: partition.clone(), - session_id: self.session_id.clone(), - worker_id: self.binding_target.worker_id, - dp_rank: self.binding_target.dp_rank, - sequence: self.version.sequence, - writer_id: self.version.writer_id, - } - } - pub fn session_id(&self) -> &str { &self.session_id } @@ -1029,7 +965,6 @@ impl AffinityLease { let next_version = inner.next_version(); *current = target; *version = next_version; - self.binding_target = target; self.version = next_version; true } @@ -1078,9 +1013,7 @@ impl AffinityLease { *idle_deadline = inner.now() + inner.ttl; (*target, *version) }; - if self.publish_binding_on_release { - inner.publish_replica_update(&self.session_id, target, version); - } + inner.publish_replica_update(&self.session_id, target, version); } /// Drop the binding this lease holds (the bound worker is gone). A binding diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 1ff6688d438f..f5621c83d6d7 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -14,7 +14,7 @@ use std::time::{Duration, Instant}; use dynamo_tokens::SequenceHash; use once_cell::sync::OnceCell; use parking_lot::RwLock; -use tokio::sync::watch; +use tokio::sync::{mpsc, watch}; use tokio_util::sync::CancellationToken; use tracing::Instrument; @@ -42,8 +42,9 @@ use crate::sequences::{ ReplicaWorkerPolicy, SequenceRequest, SequenceTrackerOptions, active_request_expiry_duration, }; use crate::services::common::replica_sync::{ - HostReplicaSyncFactory, ReplicaInbox, ReplicaSyncConfig, SchedulerLoadSink, ScopedReplicaEvent, - ScopedSequencePublisher, setup_scoped_replica_sync, + AFFINITY_EVENT_CHANNEL_CAPACITY, AffinityMeshSink, HostReplicaSyncFactory, ReplicaInbox, + ReplicaSyncConfig, SchedulerLoadSink, ScopedReplicaEvent, ScopedSequencePublisher, + setup_scoped_replica_sync, }; use crate::services::indexer::backend::{Indexer, IndexerPolicy}; use crate::services::indexer::recovery; @@ -129,6 +130,7 @@ struct SelectionEntry { replica_inbox: Option, affinity: OnceCell, replica_config: Option, + affinity_replication: Option<(u64, Arc)>, } impl SelectionEntry { @@ -144,6 +146,8 @@ impl SelectionEntry { && let Some(sink) = config.affinity_sink(&self.key) { table.enable_replication(config.process_id(), sink); + } else if let Some((writer_id, sink)) = &self.affinity_replication { + table.enable_replication(*writer_id, sink.clone()); } Ok(table) })?; @@ -224,9 +228,13 @@ pub struct HostTelemetry { } /// Replica-sync transport the host carries for partitions this core does not -/// mesh itself (ignored when the service runs its own ZMQ replica sync). +/// mesh itself. Load channels are ignored when the service runs its own ZMQ +/// replica sync; session bindings cannot be combined with it. #[derive(Clone)] pub struct HostReplication { + /// Identity and publisher for host-carried session bindings. + pub writer_id: u64, + pub affinity_tx: Option>, pub channels: Option, /// Owns request expiry when supplied; the partition then expires only through lifecycle events. pub request_leases: Option>, @@ -238,6 +246,8 @@ pub struct HostReplication { impl Default for HostReplication { fn default() -> Self { Self { + writer_id: 0, + affinity_tx: None, channels: None, request_leases: None, replica_worker_policy: ReplicaWorkerPolicy::RequireRegistered, @@ -245,6 +255,36 @@ impl Default for HostReplication { } } +impl HostReplication { + pub fn affinity_channel(writer_id: u64) -> (Self, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(AFFINITY_EVENT_CHANNEL_CAPACITY); + ( + Self { + writer_id, + affinity_tx: Some(tx), + ..Self::default() + }, + rx, + ) + } + + fn affinity_sink( + &self, + partition: &RoutingPartitionId, + ) -> Option<(u64, Arc)> { + self.affinity_tx.clone().map(|tx| { + let sink = AffinityMeshSink { + partition: partition.clone(), + tx, + }; + ( + self.writer_id, + Arc::new(sink) as Arc, + ) + }) + } +} + #[derive(Debug, Clone)] pub struct SelectionServiceConfig { pub port: u16, @@ -418,7 +458,7 @@ impl SelectionCore { self.cancel_token.cancel(); } - fn ensure_running(&self) -> Result<(), SelectionError> { + pub(super) fn ensure_running(&self) -> Result<(), SelectionError> { if self.cancel_token.is_cancelled() { return Err(SelectionError::NotReady( "selection service is shutting down".to_string(), diff --git a/lib/kv-router/src/services/selection/core/operation.rs b/lib/kv-router/src/services/selection/core/operation.rs index fff971a6712c..2817ee94f1ef 100644 --- a/lib/kv-router/src/services/selection/core/operation.rs +++ b/lib/kv-router/src/services/selection/core/operation.rs @@ -21,7 +21,6 @@ use crate::scheduling::{AdvisoryWorkerLoad, QueueRejection, SchedulingResponse, use super::super::error::SelectionError; use super::super::input::PromptView; -use crate::services::common::replica_sync::AffinityBindingEvent; /// Every input to one selection. Fields are independent: `session_context` /// feeds worker selection, `session` says what the core does with the session @@ -139,8 +138,6 @@ pub struct Selected { pub sequence_hashes: Option>, pub track_prefill_tokens: bool, pub effective_prefill_tokens: usize, - /// The versioned session assignment committed by this booking, if any. - pub session_binding: Option, pub kv_hint: Option, pub routing_hashes: Option>, pub shared_cache_hits: Option, diff --git a/lib/kv-router/src/services/selection/core/reservations.rs b/lib/kv-router/src/services/selection/core/reservations.rs index 8e7b19acf8b5..f12c59ad9e33 100644 --- a/lib/kv-router/src/services/selection/core/reservations.rs +++ b/lib/kv-router/src/services/selection/core/reservations.rs @@ -25,7 +25,6 @@ struct ReservationBooking { /// Public block hashes to record into an approximate indexer once booked. routing_hashes: Option>, session_id: Option, - affinity_lease: Option, } /// `selection_id` -> the booking it holds. @@ -172,52 +171,21 @@ fn forget_reservation_if( impl SelectionCore { pub async fn create_reservation( &self, - mut req: ReservationRequest, - ) -> Result, SelectionError> { + req: ReservationRequest, + ) -> Result { self.ensure_running()?; - // Apply the peer's session affinity even if load booking is skipped or fails. - // Successful bookings retain a lease to keep affinity alive until completion. - let affinity_lease = if let Some(binding) = req.session_binding.take() { - let key = RoutingPartitionId::new(&req.model_name, &req.routing_group); - if binding.partition != key { - return Err(SelectionError::BadRequest( - "session binding must belong to the reservation's routing partition" - .to_string(), - )); - } - let session_id = binding.session_id.clone(); - let target = binding.target(); - let version = binding.version(); - self.dispatch_affinity_event(binding); - if req.skip_load_booking { - None - } else { - self.entry(&key).and_then(|entry| { - entry - .affinity - .get()? - .acquire_mirrored_reservation_lease(session_id, target, version) - }) - } - } else { - None - }; - if req.skip_load_booking { - return Ok(None); - } tokio::select! { biased; _ = self.cancel_token.cancelled() => { Err(SelectionError::Scheduler(KvSchedulerError::SubscriberShutdown)) } - result = self.create_reservation_inner(req, affinity_lease) => result.map(Some), + result = self.create_reservation_inner(req) => result, } } async fn create_reservation_inner( &self, req: ReservationRequest, - affinity_lease: Option, ) -> Result { let key = RoutingPartitionId::new(req.model_name.clone(), req.routing_group.clone()); @@ -225,9 +193,7 @@ impl SelectionCore { // any cached selection for the id so a later replay can't book stale state. if let Some(worker_id) = req.worker_id { self.selection_cache.discard(&key, &req.selection_id); - return self - .reserve_explicit(key, worker_id, req, affinity_lease) - .await; + return self.reserve_explicit(key, worker_id, req).await; } // Replay form: peek, book, and consume only once the booking lands. A @@ -243,10 +209,7 @@ impl SelectionCore { req.selection_id ))); }; - let response = match self - .book_cached_selection(pending, &req, affinity_lease) - .await - { + let response = match self.book_cached_selection(pending, &req).await { Ok(response) => response, // The session moved: the cached worker can never be booked for it. Err(SelectionError::BadRequest(message)) => { @@ -267,7 +230,6 @@ impl SelectionCore { &self, pending: Arc, req: &ReservationRequest, - affinity_lease: Option, ) -> Result { let (entry, endpoint, prefill_load_hint) = self.resolve_cached_booking(&pending)?; let track_prefill_tokens = pending.track_prefill_tokens; @@ -285,7 +247,6 @@ impl SelectionCore { lora_name: pending.lora_name.clone(), routing_hashes: pending.routing_hashes.clone(), session_id: pending.session_id.clone(), - affinity_lease, }, ) .await @@ -338,7 +299,6 @@ impl SelectionCore { key: RoutingPartitionId, worker_id: WorkerId, req: ReservationRequest, - affinity_lease: Option, ) -> Result { let entry = self.ready_entry(&key)?; let normalized = req.prompt.view().normalize_for_reservation( @@ -395,7 +355,6 @@ impl SelectionCore { lora_name: req.prompt.lora_name, routing_hashes, session_id: None, - affinity_lease, }, ) .await @@ -421,7 +380,6 @@ impl SelectionCore { lora_name, routing_hashes, session_id, - affinity_lease, } = booking; let claim = self.claim_reservation(&selection_id, &key)?; @@ -452,7 +410,7 @@ impl SelectionCore { Some((table, session_id, hold)) => { self.commit_session(table, hold, session_id, worker, &key)? } - None => affinity_lease, + None => None, }; if let Some(hashes) = routing_hashes { self.record_routing_decision(&entry, worker, hashes).await; diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 9937c63dbd27..064bd0764c10 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -195,7 +195,6 @@ impl SelectionCore { sequence_hashes, track_prefill_tokens, effective_prefill_tokens, - session_binding, kv_hint, routing_hashes: _, shared_cache_hits: _, @@ -220,7 +219,6 @@ impl SelectionCore { .map(|threshold| load.prefill_load_exceeds(threshold)), }); SelectResponse { - session_binding, selection_id, sequence_hashes: sequence_hashes .map(|hashes| hashes.into_iter().map(|hash| hash as i64).collect()), @@ -509,7 +507,6 @@ impl SelectionCore { // The routing hashes go to exactly one of: the reservation recorded // now, or the replay cache a later reservation records from. let mut routing_hashes = routing_hashes; - let mut session_binding = None; let booking = if let Some(claim) = claim { let Some(booking) = booking else { return Err(SelectionError::Internal( @@ -524,7 +521,6 @@ impl SelectionCore { } _ => None, }; - session_binding = affinity_lease.as_ref().map(|lease| lease.binding(&key)); if let Some(hashes) = routing_hashes.take() { self.record_routing_decision(&entry, response.best_worker, hashes) .await; @@ -566,7 +562,6 @@ impl SelectionCore { sequence_hashes: booked_sequence_hashes, track_prefill_tokens, effective_prefill_tokens: effective_prefill, - session_binding, kv_hint, routing_hashes: returned_routing_hashes, shared_cache_hits: host_shared_cache_hits, @@ -589,13 +584,14 @@ impl SelectionCore { match table.acquire(session_id, None).await { Ok(Hold::Bound { target, mut lease }) if !self.catalog.is_schedulable(target, key) - || allowed_worker_ids - .is_some_and(|ids| !ids.contains(&target.worker_id)) => + || (table.mode() == SessionAffinityMode::Soft + && allowed_worker_ids + .is_some_and(|ids| !ids.contains(&target.worker_id))) => { tracing::debug!( session_id, worker_id = target.worker_id, - "Session affinity target is not schedulable; re-initializing" + "Session affinity target is unavailable or excluded; re-initializing" ); lease.invalidate(); drop(lease); diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 49d89d35fe5c..38828b398538 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -111,16 +111,6 @@ fn replay_reservation(selection_id: &str) -> ReservationRequest { expected_output_tokens: None, effective_prefill_tokens: None, track_prefill_tokens: None, - session_binding: None, - skip_load_booking: false, - } -} - -fn peer_affinity_update(binding: AffinityBindingEvent) -> ReservationRequest { - ReservationRequest { - session_binding: Some(binding), - skip_load_booking: true, - ..replay_reservation("peer-update") } } @@ -3057,8 +3047,7 @@ async fn failed_replay_preserves_a_newer_cached_selection() { let replacement = core .create_reservation(replay_reservation("pending")) .await - .expect("new selection survives the old replay failure") - .expect("reservation booked"); + .expect("new selection survives the old replay failure"); assert_eq!(replacement.worker_id, 2); } @@ -3103,24 +3092,49 @@ async fn free_releases_reservation_and_affinity_without_actor_progress() { .expect("completed cleanup must not prevent ID reuse"); } -#[rstest::rstest] -#[case::hard(SessionAffinityMode::Hard)] -#[case::soft(SessionAffinityMode::Soft)] #[tokio::test] -async fn session_rebinds_outside_allowlist(#[case] mode: SessionAffinityMode) { - let (core, first) = bound_session(mode).await; +async fn hard_mode_rejects_dispatch_away_from_a_live_binding() { + let (core, first) = bound_session(SessionAffinityMode::Hard).await; + let other = if first.worker_id == 1 { 2 } else { 1 }; + + // Steering cannot reach the bound worker, so selection lands elsewhere. + let mut request = session_reservation("r2", "s"); + request.allowed_worker_ids = Some(HashSet::from([other])); + let err = core + .select_and_reserve(request) + .await + .expect_err("hard affinity rejects a dispatch away from the binding"); + assert!(matches!(err, SelectionError::BadRequest(_)), "{err:?}"); + let entry = core.entry(&default_key()).expect("entry"); + wait_until("rejected booking release", || { + !entry.scheduler.has_request("r2") + }) + .await; + assert!(core.reservation_index.read().is_empty()); + assert_eq!(bound_worker(&core, "s"), None, "stale binding is dropped"); + + // The session's next request re-initializes. + let mut request = session_reservation("r3", "s"); + request.allowed_worker_ids = Some(HashSet::from([other])); + let third = core.select_and_reserve(request).await.expect("rebind"); + assert_eq!(third.worker_id, other); + assert_eq!(bound_worker(&core, "s"), Some(other)); +} + +#[tokio::test] +async fn soft_mode_follows_the_dispatch() { + let (core, first) = bound_session(SessionAffinityMode::Soft).await; let other = if first.worker_id == 1 { 2 } else { 1 }; - // Excluding the bound worker permits failover in either affinity mode. + let mut request = session_reservation("r2", "s"); request.allowed_worker_ids = Some(HashSet::from([other])); - let second = core.select_and_reserve(request).await.expect("rebind"); + let second = core.select_and_reserve(request).await.expect("soft rebind"); assert_eq!(second.worker_id, other); assert_eq!(bound_worker(&core, "s"), Some(other)); - assert_eq!(lease_count(&core, "s"), Some(1)); } #[tokio::test] -async fn hard_affinity_rejects_policy_mismatch() { +async fn soft_binding_follows_policy_override() { // Bind to an eligible worker different from the custom policy's choice. let (policy, _) = capturing_policy_factory(); let core = core_with( @@ -3128,7 +3142,10 @@ async fn hard_affinity_rejects_policy_mismatch() { SelectionHost::default(), Some(policy), WorkerType::Aggregated, - Some(SessionAffinityConfig::new(Duration::from_secs(10))), + Some( + SessionAffinityConfig::new(Duration::from_secs(10)) + .with_mode(SessionAffinityMode::Soft), + ), ); for id in [1, 2] { core.upsert_worker(worker(id)).await.unwrap(); @@ -3144,352 +3161,14 @@ async fn hard_affinity_rejects_policy_mismatch() { writer_id: 99, }); assert_eq!(bound_worker(&core, "s"), Some(bound)); - // Hard affinity rejects the dispatch and releases its booking and binding. - let result = core - .select_and_reserve(session_reservation("next-turn", "s")) - .await; - assert!(matches!(result, Err(SelectionError::BadRequest(_)))); - let entry = core.entry(&default_key()).unwrap(); - wait_until("rejected booking release", || { - !entry.scheduler.has_request("next-turn") - }) - .await; - assert!(core.reservation_index.read().is_empty()); - assert_eq!(bound_worker(&core, "s"), None); - - // The next turn can initialize a fresh binding after the rejection. let next = core - .select_and_reserve(session_reservation("retry", "s")) + .select_and_reserve(session_reservation("next-turn", "s")) .await - .expect("rebind after rejection"); + .unwrap(); assert_eq!(next.worker_id, preferred); assert_eq!(bound_worker(&core, "s"), Some(preferred)); } -struct CountingAffinitySink(AtomicUsize); - -impl super::super::affinity::AffinityReplicaSink for CountingAffinitySink { - fn publish( - &self, - _: &str, - _: WorkerAffinityTarget, - _: super::super::affinity::AffinityVersion, - ) { - self.0.fetch_add(1, Ordering::Relaxed); - } -} - -#[tokio::test(start_paused = true)] -async fn binding_only_update_preserves_ttl() { - // Import an idle binding with a sink that detects accidental republication. - let (source, first) = bound_session(SessionAffinityMode::Hard).await; - let peer = core_with_session_affinity(); - peer.upsert_worker(worker(first.worker_id)).await.unwrap(); - let key = default_key(); - let publications = Arc::new(CountingAffinitySink(AtomicUsize::new(0))); - let entry = peer.entry(&key).unwrap(); - assert!( - entry - .affinity - .get() - .unwrap() - .enable_replication(99, publications.clone()) - ); - let loads_before = serde_json::to_value(peer.loads(None, None)).unwrap(); - let event = first.session_binding.unwrap(); - assert_eq!(event.worker_id, first.worker_id); - // Skipping load booking must leave the peer's loads unchanged. - assert!( - peer.create_reservation(peer_affinity_update(event)) - .await - .unwrap() - .is_none() - ); - assert_eq!( - serde_json::to_value(peer.loads(None, None)).unwrap(), - loads_before - ); - // Both idle bindings expire at the configured TTL without publishing. - tokio::time::advance(Duration::from_secs(9)).await; - assert!(bound_worker(&source, "s").is_some()); - assert!(bound_worker(&peer, "s").is_some()); - tokio::time::advance(Duration::from_secs(2)).await; - assert!(bound_worker(&source, "s").is_none()); - assert!(bound_worker(&peer, "s").is_none()); - assert_eq!(publications.0.load(Ordering::Relaxed), 0); -} - -#[tokio::test(start_paused = true)] -async fn mirrored_lease_preserves_ttl() { - // Mirror the source reservation and retain a lease on both services. - let source = core_with_session_affinity(); - let peer = core_with_session_affinity(); - for core in [&source, &peer] { - core.upsert_worker(worker(1)).await.unwrap(); - } - let key = default_key(); - let publications = Arc::new(CountingAffinitySink(AtomicUsize::new(0))); - let entry = peer.entry(&key).unwrap(); - assert!( - entry - .affinity - .get() - .unwrap() - .enable_replication(99, publications.clone()) - ); - - let selected = source - .select_and_reserve(session_reservation("turn", "s")) - .await - .unwrap(); - let binding = selected.session_binding.unwrap(); - peer.create_reservation(ReservationRequest { - worker_id: Some(selected.worker_id), - dp_rank: Some(selected.dp_rank), - prompt: prompt(), - session_binding: Some(binding.clone()), - ..replay_reservation("turn") - }) - .await - .unwrap() - .expect("mirrored booking"); - for core in [&source, &peer] { - assert_eq!(lease_count(core, "s"), Some(1)); - } - let mirrored = peer - .reservation_index - .read() - .get("turn") - .unwrap() - ._affinity_lease - .as_ref() - .unwrap() - .binding(&key); - assert_eq!(mirrored.version(), binding.version()); - - // A running request keeps both bindings alive beyond the idle TTL. - tokio::time::advance(Duration::from_secs(11)).await; - for core in [&source, &peer] { - assert_eq!( - bound_worker(core, "s"), - Some(1), - "an active turn cannot expire" - ); - core.free_reservation("turn").await.unwrap(); - assert_eq!(lease_count(core, "s"), Some(0)); - } - // Completion starts a full idle TTL; the peer's lease must not republish. - tokio::time::advance(Duration::from_secs(9)).await; - for core in [&source, &peer] { - assert_eq!( - bound_worker(core, "s"), - Some(1), - "idle TTL starts at completion" - ); - } - tokio::time::advance(Duration::from_secs(2)).await; - for core in [&source, &peer] { - assert_eq!(bound_worker(core, "s"), None); - } - assert_eq!(publications.0.load(Ordering::Relaxed), 0); -} - -#[tokio::test] -async fn peer_bindings_converge() { - // Create competing assignments from independent services. - let source = core_with_session_affinity(); - let peer = core_with_session_affinity(); - for core in [&source, &peer] { - for id in [1, 2, 3] { - core.upsert_worker(worker(id)).await.unwrap(); - } - } - let mut bindings = Vec::new(); - for (core, worker_id) in [(&source, 1), (&peer, 2)] { - let mut request = session_reservation("first", "s"); - request.allowed_worker_ids = Some(HashSet::from([worker_id])); - let selected = core.select_and_reserve(request).await.unwrap(); - bindings.push(selected.session_binding.unwrap()); - } - let mut first = bindings.remove(0); - let mut concurrent = bindings.remove(0); - // Exercise a sequence tie between independent writers. - let sequence = first.sequence.max(concurrent.sequence); - first.sequence = sequence; - concurrent.sequence = sequence; - assert_ne!(first.writer_id, concurrent.writer_id); - let winner = if first.version() > concurrent.version() { - first.clone() - } else { - concurrent.clone() - }; - // Opposite delivery orders must choose the same version winner. - for (core, events) in [ - (&source, [&first, &concurrent]), - (&peer, [&concurrent, &first]), - ] { - for event in events { - core.create_reservation(peer_affinity_update(event.clone())) - .await - .unwrap(); - } - } - for core in [&source, &peer] { - assert_eq!(bound_worker(core, "s"), Some(winner.worker_id)); - } - - // A later rebind wins even when older updates are replayed afterward. - let mut request = session_reservation("rebind", "s"); - request.allowed_worker_ids = Some(HashSet::from([3])); - let selected = source.select_and_reserve(request).await.unwrap(); - let newer = selected.session_binding.unwrap(); - assert!(newer.version() > winner.version()); - peer.create_reservation(peer_affinity_update(newer)) - .await - .unwrap(); - for core in [&source, &peer] { - for event in [&first, &concurrent] { - core.create_reservation(peer_affinity_update(event.clone())) - .await - .unwrap(); - } - assert_eq!(bound_worker(core, "s"), Some(3)); - } -} - -#[rstest::rstest] -#[case::book(false, false)] -#[case::skip(true, false)] -#[case::fail(false, true)] -#[tokio::test] -async fn affinity_survives_booking_outcomes( - #[case] skip_load_booking: bool, - #[case] invalid_rank: bool, -) { - // Import the binding with successful, skipped, or failed load booking. - let (_, first) = bound_session(SessionAffinityMode::Hard).await; - let peer = core_with_session_affinity(); - for id in [1, 2] { - peer.upsert_worker(worker(id)).await.unwrap(); - } - let binding = first.session_binding.unwrap(); - let result = peer - .create_reservation(ReservationRequest { - worker_id: Some(first.worker_id), - dp_rank: Some(if invalid_rank { 99 } else { 0 }), - prompt: prompt(), - session_binding: Some(binding.clone()), - skip_load_booking, - ..replay_reservation("mirror") - }) - .await; - if invalid_rank { - assert!(result.is_err()); - } else { - assert_eq!(result.unwrap().is_none(), skip_load_booking); - } - // Every outcome imports affinity; only successful bookings retain load/leases. - assert_eq!(bound_worker(&peer, "s"), Some(first.worker_id)); - assert_eq!( - lease_count(&peer, "s"), - Some(usize::from(!skip_load_booking && !invalid_rank)) - ); - let active: usize = peer - .loads(None, None) - .iter() - .flat_map(|partition| &partition.loads) - .map(|load| load.active_requests) - .sum(); - assert_eq!(active, usize::from(!skip_load_booking && !invalid_rank)); - // The next turn reuses the imported assignment and version. - let next = peer - .select_and_reserve(session_reservation("next", "s")) - .await - .unwrap(); - let next_binding = next.session_binding.unwrap(); - assert_eq!(next_binding.worker_id, binding.worker_id); - assert_eq!(next_binding.version(), binding.version()); -} - -#[tokio::test] -async fn binding_rejects_wrong_partition() { - let (core, first) = bound_session(SessionAffinityMode::Hard).await; - // Attach another model's binding to this partition's reservation. - let mut binding = first.session_binding.unwrap(); - binding.partition = RoutingPartitionId::new("other", "default"); - let request = ReservationRequest { - worker_id: Some(first.worker_id), - dp_rank: Some(first.dp_rank), - prompt: prompt(), - session_binding: Some(binding), - ..replay_reservation("mismatch") - }; - // Reject the mismatch before creating a load booking. - assert!(matches!( - core.create_reservation(request).await, - Err(SelectionError::BadRequest(_)) - )); - assert!(core.reservation_index.read().is_empty()); - assert!( - !core - .entry(&default_key()) - .unwrap() - .scheduler - .has_request("mismatch") - ); -} - -#[rstest::rstest] -#[case::different_worker(false)] -#[case::same_worker(true)] -#[tokio::test] -async fn stale_binding_preserves_load(#[case] return_to_original_worker: bool) { - // Rebind an active session, optionally returning to its original worker. - let source = core_with_session_affinity(); - let peer = core_with_session_affinity(); - for core in [&source, &peer] { - for id in [1, 2] { - core.upsert_worker(worker(id)).await.unwrap(); - } - } - let mut old_request = session_reservation("old", "s"); - old_request.allowed_worker_ids = Some(HashSet::from([1])); - let old = source.select_and_reserve(old_request).await.unwrap(); - let mut new_request = session_reservation("new", "s"); - new_request.allowed_worker_ids = Some(HashSet::from([2])); - let mut new = source.select_and_reserve(new_request).await.unwrap(); - if return_to_original_worker { - let mut returning = session_reservation("returning", "s"); - returning.allowed_worker_ids = Some(HashSet::from([1])); - new = source.select_and_reserve(returning).await.unwrap(); - } - let replacement = new.worker_id; - // Deliver the newest binding before the old request's delayed reservation. - peer.create_reservation(peer_affinity_update(new.session_binding.unwrap())) - .await - .unwrap(); - peer.create_reservation(ReservationRequest { - worker_id: Some(old.worker_id), - dp_rank: Some(old.dp_rank), - prompt: prompt(), - session_binding: old.session_binding, - ..replay_reservation("old") - }) - .await - .unwrap() - .expect("old request still needs load accounting"); - // Count the original request without reverting or leasing the newer binding. - assert_eq!(bound_worker(&peer, "s"), Some(replacement)); - let loads = peer.loads(None, None); - let original_worker = loads - .iter() - .flat_map(|p| &p.loads) - .find(|w| w.worker_id == 1) - .unwrap(); - assert_eq!(original_worker.active_requests, 1); - assert_eq!(lease_count(&peer, "s"), Some(0)); -} - #[tokio::test] async fn replicated_binding_steers_a_new_session_and_frees_with_the_booking() { let core = core_with_session_affinity(); @@ -4281,3 +3960,107 @@ async fn overflowing_worker_dp_range_patch_preserves_live_worker() { .await .expect("invalid worker patch deadline"); } + +fn core_with_binding_stream( + writer_id: u64, +) -> (SelectionCore, mpsc::Receiver) { + let (replication, rx) = HostReplication::affinity_channel(writer_id); + let core = core_with( + test_config(false), + SelectionHost { + replication, + ..SelectionHost::default() + }, + None, + WorkerType::Aggregated, + Some( + SessionAffinityConfig::new(Duration::from_secs(10)) + .with_mode(SessionAffinityMode::Soft), + ), + ); + (core, rx) +} + +#[tokio::test(start_paused = true)] +async fn host_stream_publishes_bindings_and_refreshes() { + let (source, mut updates) = core_with_binding_stream(11); + let (peer, mut peer_updates) = core_with_binding_stream(22); + for core in [&source, &peer] { + core.upsert_worker(worker(1)).await.unwrap(); + } + source + .select_and_reserve(session_reservation("turn", "s")) + .await + .unwrap(); + let binding = updates.try_recv().expect("commit publishes synchronously"); + assert_eq!(binding.partition, default_key()); + assert_eq!(binding.writer_id, 11); + assert_eq!(binding.worker_id, 1); + peer.dispatch_affinity_event(binding.clone()); + assert_eq!(bound_worker(&peer, "s"), Some(1)); + assert!(peer_updates.try_recv().is_err(), "imports must not echo"); + assert!(peer.reservation_index.read().is_empty()); + assert_eq!(lease_count(&peer, "s"), Some(0)); + + // Shared frontend semantics: only the origin has an active lease. Release + // restores a peer that expired during a request longer than the idle TTL. + tokio::time::advance(Duration::from_secs(11)).await; + assert_eq!(bound_worker(&source, "s"), Some(1)); + assert_eq!(bound_worker(&peer, "s"), None); + source.free_reservation("turn").await.unwrap(); + let refresh = updates.try_recv().unwrap(); + assert_eq!(refresh.version(), binding.version()); + peer.dispatch_affinity_event(refresh); + tokio::time::advance(Duration::from_secs(9)).await; + assert_eq!(bound_worker(&source, "s"), Some(1)); + assert_eq!(bound_worker(&peer, "s"), Some(1)); + tokio::time::advance(Duration::from_secs(2)).await; + assert_eq!(bound_worker(&source, "s"), None); + assert_eq!(bound_worker(&peer, "s"), None); + assert!(peer_updates.try_recv().is_err()); +} + +#[tokio::test] +async fn peer_bindings_converge_by_version() { + let (source, mut source_updates) = core_with_binding_stream(11); + let (peer, mut peer_updates) = core_with_binding_stream(22); + for core in [&source, &peer] { + for id in [1, 2, 3] { + core.upsert_worker(worker(id)).await.unwrap(); + } + } + for (core, id) in [(&source, 1), (&peer, 2)] { + let mut request = session_reservation("first", "s"); + request.allowed_worker_ids = Some(HashSet::from([id])); + core.select_and_reserve(request).await.unwrap(); + } + let mut first = source_updates.try_recv().unwrap(); + let mut concurrent = peer_updates.try_recv().unwrap(); + let sequence = first.sequence.max(concurrent.sequence); + first.sequence = sequence; + concurrent.sequence = sequence; + for (core, events) in [ + (&source, [&first, &concurrent]), + (&peer, [&concurrent, &first]), + ] { + for event in events { + core.dispatch_affinity_event(event.clone()); + } + assert_eq!( + bound_worker(core, "s"), + Some(2), + "host writer id breaks sequence ties" + ); + } + let mut request = session_reservation("rebind", "s"); + request.allowed_worker_ids = Some(HashSet::from([3])); + source.select_and_reserve(request).await.unwrap(); + let newer = source_updates.try_recv().unwrap(); + assert!(newer.version() > concurrent.version()); + peer.dispatch_affinity_event(newer); + peer.free_reservation("first").await.unwrap(); + assert_eq!(bound_worker(&peer, "s"), Some(3)); + // A stale version cannot displace the newer binding. + peer.dispatch_affinity_event(first); + assert_eq!(bound_worker(&peer, "s"), Some(3)); +} diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index 5d6e9f452781..9326bfc648aa 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -374,6 +374,7 @@ impl SelectionCore { replica_inbox, affinity: OnceCell::new(), replica_config: self.replica_config.clone(), + affinity_replication: self.host.replication.affinity_sink(&key), })) })? .clone(); diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index 27467ae69b67..7ee1fdb042a1 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -30,8 +30,8 @@ pub use crate::plugins::{ WorkerSelectionPolicyRegistryError, }; pub use crate::services::common::replica_sync::{ - HostReplicaChannels, HostReplicaSyncFactory, ReplicaIngressObserver, ReplicaPeerError, - SchedulerLoadSink, + AffinityBindingEvent, HostReplicaChannels, HostReplicaSyncFactory, ReplicaIngressObserver, + ReplicaPeerError, SchedulerLoadSink, }; pub use core::{ HostCache, HostEligibility, HostLoad, HostReplication, HostTelemetry, KvIndexSource, Selected, @@ -45,7 +45,8 @@ pub use membership::{CatalogObserver, CatalogReconciler, WorkerCatalogSource}; pub use pending::SelectionCacheConfig; pub use server::{AppState, run_server}; pub use service::{ - SelectionService, SelectionServiceBuilder, warn_for_unserved_worker_selection_policies, + SelectionService, SelectionServiceBuilder, SessionBindingUpdates, + warn_for_unserved_worker_selection_policies, }; pub use types::{ DEFAULT_MODEL_NAME, ModelLoadResponse, OutputBlockRequest, OverlapScoresRequest, diff --git a/lib/kv-router/src/services/selection/server.rs b/lib/kv-router/src/services/selection/server.rs index fa9e4a8d1360..cfb6af7ab39f 100644 --- a/lib/kv-router/src/services/selection/server.rs +++ b/lib/kv-router/src/services/selection/server.rs @@ -129,8 +129,7 @@ async fn create_reservation( Err(error) => return json_rejection(error), }; match state.service.create_reservation(req).await { - Ok(Some(response)) => (StatusCode::CREATED, Json(response)).into_response(), - Ok(None) => StatusCode::NO_CONTENT.into_response(), + Ok(response) => (StatusCode::CREATED, Json(response)).into_response(), Err(error) => error.into_response(), } } diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 0fb891028a98..1ead5ba38494 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -4,6 +4,8 @@ use std::sync::Arc; use std::time::Duration; +use parking_lot::Mutex; +use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use crate::WorkerSelectionPolicyFactory; @@ -11,13 +13,15 @@ use crate::config::KvRouterConfig; use crate::protocols::WorkerId; use crate::scheduling::PotentialLoad; use crate::services::common::replica_sync::{ - PeerManager, ReplicaPeerError, ReplicaSyncRuntime, setup_replica_sync, + AffinityBindingEvent, PeerManager, ReplicaPeerError, ReplicaSyncRuntime, setup_replica_sync, }; use crate::services::indexer::backend::IndexerPolicy; use crate::tracking_hash::TrackingHashContext; -use super::affinity::SessionAffinityConfig; -use super::core::{KvIndexSource, SelectionCore, SelectionHost, SelectionServiceConfig}; +use super::affinity::{SessionAffinityConfig, SessionAffinityMode}; +use super::core::{ + HostReplication, KvIndexSource, SelectionCore, SelectionHost, SelectionServiceConfig, +}; use super::error::SelectionError; use super::pending::SelectionCacheConfig; use super::types::{ @@ -40,6 +44,8 @@ pub struct SelectionServiceBuilder { host: SelectionHost, worker_selection_policy_factory: Option, session_affinity_ttl: Option, + session_affinity_mode: SessionAffinityMode, + session_replication_writer_id: Option, host_manages_request_lifecycle: bool, } @@ -81,6 +87,8 @@ impl SelectionServiceBuilder { host: SelectionHost::default(), worker_selection_policy_factory: None, session_affinity_ttl: None, + session_affinity_mode: SessionAffinityMode::default(), + session_replication_writer_id: None, host_manages_request_lifecycle: false, } } @@ -125,6 +133,20 @@ impl SelectionServiceBuilder { self } + /// Choose whether a live binding must be honored or may follow a new dispatch. + pub fn session_affinity_mode(mut self, mode: SessionAffinityMode) -> Self { + self.session_affinity_mode = mode; + self + } + + /// Export bindings through a bounded channel instead of the built-in mesh. + /// The host supplies an id unique to this replica incarnation and carries + /// both commit updates and release-time TTL refreshes to its peers. + pub fn session_binding_replication(mut self, writer_id: u64) -> Self { + self.session_replication_writer_id = Some(writer_id); + self + } + pub fn replica_sync(mut self, port: u16, peers: Vec) -> Self { self.replica_sync_port = Some(port); self.replica_sync_peers = peers; @@ -146,6 +168,26 @@ impl SelectionServiceBuilder { } pub async fn build(mut self) -> anyhow::Result { + let session_binding_rx = if let Some(writer_id) = self.session_replication_writer_id { + anyhow::ensure!( + self.session_affinity_ttl.is_some(), + "session binding replication requires session affinity" + ); + anyhow::ensure!( + self.host.replication.affinity_tx.is_none(), + "session binding publisher already configured" + ); + let (replication, rx) = HostReplication::affinity_channel(writer_id); + self.host.replication.writer_id = replication.writer_id; + self.host.replication.affinity_tx = replication.affinity_tx; + Some(Mutex::new(Some(rx))) + } else { + None + }; + anyhow::ensure!( + self.replica_sync_port.is_none() || self.host.replication.affinity_tx.is_none(), + "host session binding replication cannot be combined with replica_sync" + ); self.kv_router_config .apply_policy_config() .map_err(anyhow::Error::msg)?; @@ -197,7 +239,8 @@ impl SelectionServiceBuilder { self.selection_cache, tracking_hash, indexer_policy, - self.session_affinity_ttl.map(SessionAffinityConfig::new), + self.session_affinity_ttl + .map(|ttl| SessionAffinityConfig::new(ttl).with_mode(self.session_affinity_mode)), )); if recover_from_peers { @@ -241,6 +284,7 @@ impl SelectionServiceBuilder { startup_guard.disarm(); Ok(SelectionService { core, + session_binding_rx, peer_manager, replica_runtime, replica_sync_port: self.replica_sync_port, @@ -301,13 +345,64 @@ impl Drop for StartupGuard { pub struct SelectionService { core: Arc, + /// `None` when replication is disabled; the inner `None` once taken. + session_binding_rx: Option>>>, peer_manager: Option>, replica_runtime: Option, replica_sync_port: Option, cancel_token: CancellationToken, } +/// Session bindings that terminate when the service shuts down. +pub struct SessionBindingUpdates { + rx: tokio::sync::Mutex>, + cancel_token: CancellationToken, +} + +impl SessionBindingUpdates { + pub async fn recv(&self) -> Option { + let mut rx = self.rx.lock().await; + tokio::select! { + biased; + _ = self.cancel_token.cancelled() => None, + event = rx.recv() => event, + } + } +} + impl SelectionService { + /// Session bindings this replica commits or refreshes, for the host to + /// broadcast to its peers. Ends when the service shuts down. + pub fn session_binding_updates(&self) -> Result { + let rx = self + .session_binding_rx + .as_ref() + .ok_or_else(|| { + SelectionError::Conflict("session binding stream is disabled".to_string()) + })? + .lock() + .take() + .ok_or_else(|| { + SelectionError::Conflict("session binding stream already taken".to_string()) + })?; + Ok(SessionBindingUpdates { + rx: tokio::sync::Mutex::new(rx), + cancel_token: self.cancel_token.clone(), + }) + } + + /// Reconcile a peer binding independently of request load. Importing does + /// not publish an update. A binding this replica cannot place yet is + /// dropped and converges on the owner's next refresh. + pub fn apply_session_binding( + &self, + binding: AffinityBindingEvent, + ) -> Result<(), SelectionError> { + self.core.ensure_running()?; + self.core.dispatch_affinity_event(binding); + Ok(()) + } + #[cfg(test)] pub(super) fn new_local_for_test( kv_router_config: KvRouterConfig, @@ -315,6 +410,7 @@ impl SelectionService { ) -> Self { let cancel_token = CancellationToken::new(); Self { + session_binding_rx: None, core: Arc::new( SelectionCore::try_new_local( kv_router_config, @@ -406,12 +502,10 @@ impl SelectionService { .await } - /// Apply an optional peer binding, then book load unless `skip_load_booking` is set. - /// Affinity is retained even if booking fails; skipped bookings return `None`. pub async fn create_reservation( &self, req: ReservationRequest, - ) -> Result, SelectionError> { + ) -> Result { self.core.create_reservation(req).await } diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index ce1d625918e2..a399cea4242b 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -44,6 +44,28 @@ fn app() -> Router { create_router(Arc::new(AppState { service })) } +#[tokio::test] +async fn binding_stream_rejects_invalid_config() { + let builder = || { + SelectionServiceBuilder::new( + test_config(), + crate::WorkerType::Aggregated, + RouterPluginRegistry::default(), + ) + .session_binding_replication(11) + }; + let error = builder().build().await.err().expect("affinity disabled"); + assert!(error.to_string().contains("requires session affinity")); + let error = builder() + .session_affinity(Duration::from_secs(10)) + .replica_sync(0, Vec::new()) + .build() + .await + .err() + .expect("mesh enabled"); + assert!(error.to_string().contains("cannot be combined")); +} + struct WorkerIdScorer; impl WorkerScorer for WorkerIdScorer { diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index fcffa89f8884..26a51a93e225 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -15,7 +15,6 @@ use crate::protocols::{ use crate::scheduling::config::RouterConfigOverride; pub use crate::scheduling::{OverlapScoresResponse, SharedCacheOverlapScore, WorkerOverlapScore}; use crate::scheduling::{PotentialLoad, SessionContext, WorkerSelectionInputTrigger}; -use crate::services::common::replica_sync::AffinityBindingEvent; use crate::services::overlap::MooncakeOverlapSummary; use super::input::PromptRequest; @@ -552,8 +551,7 @@ fn resolve_session_context( /// Booking request: replay the selection cached under `selection_id`, or book /// self-contained with `worker_id`. The replay books exactly what `select` captured; -/// booking fields other than the ids and model/routing-group are ignored on replay. -/// A peer's session binding is applied independently, including when booking is skipped. +/// request fields other than the ids and model/routing-group are ignored. #[derive(Debug, Deserialize)] pub struct ReservationRequest { #[serde(default = "default_model_name")] @@ -573,11 +571,6 @@ pub struct ReservationRequest { pub expected_output_tokens: Option, pub effective_prefill_tokens: Option, pub track_prefill_tokens: Option, - /// Assignment received from a peer; applied before booking, without republishing. - pub session_binding: Option, - /// Skip load booking for requests already tracked or completed locally. - #[serde(default)] - pub skip_load_booking: bool, } #[derive(Debug, Deserialize)] @@ -609,9 +602,6 @@ pub struct OverlapScoresRequest { #[derive(Debug, Serialize)] pub struct SelectResponse { - /// Assignment to replicate alongside the reservation; absent for unpinned requests. - #[serde(skip_serializing_if = "Option::is_none")] - pub session_binding: Option, #[serde(skip_serializing_if = "Option::is_none")] pub selection_id: Option, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 5885b258e6d7..5ecab3e11cff 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -285,6 +285,7 @@ impl EmbeddedSelection { channels: replica_sync, request_leases, replica_worker_policy: ReplicaWorkerPolicy::LazyRegister, + ..HostReplication::default() }, }) .build() From 14bd4d9623420cfab163f3ebc18579176b388069 Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Sun, 25 Oct 2026 00:14:05 -0700 Subject: [PATCH 3/5] refactor(kv-router): attach session binding publisher to SelectionCore and keep base soft rebind Signed-off-by: Jeffrey Wang --- .../src/services/selection/core/mod.rs | 67 +++++++++---------- .../services/selection/core/reservations.rs | 2 +- .../src/services/selection/core/run.rs | 16 ++--- .../src/services/selection/core/tests.rs | 28 +++----- .../src/services/selection/core/workers.rs | 2 +- .../src/services/selection/service.rs | 34 +++++----- lib/llm/src/kv_router/embedded.rs | 1 - 7 files changed, 59 insertions(+), 91 deletions(-) diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index f5621c83d6d7..f7863a0beee9 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -228,13 +228,9 @@ pub struct HostTelemetry { } /// Replica-sync transport the host carries for partitions this core does not -/// mesh itself. Load channels are ignored when the service runs its own ZMQ -/// replica sync; session bindings cannot be combined with it. +/// mesh itself (ignored when the service runs its own ZMQ replica sync). #[derive(Clone)] pub struct HostReplication { - /// Identity and publisher for host-carried session bindings. - pub writer_id: u64, - pub affinity_tx: Option>, pub channels: Option, /// Owns request expiry when supplied; the partition then expires only through lifecycle events. pub request_leases: Option>, @@ -246,8 +242,6 @@ pub struct HostReplication { impl Default for HostReplication { fn default() -> Self { Self { - writer_id: 0, - affinity_tx: None, channels: None, request_leases: None, replica_worker_policy: ReplicaWorkerPolicy::RequireRegistered, @@ -255,36 +249,6 @@ impl Default for HostReplication { } } -impl HostReplication { - pub fn affinity_channel(writer_id: u64) -> (Self, mpsc::Receiver) { - let (tx, rx) = mpsc::channel(AFFINITY_EVENT_CHANNEL_CAPACITY); - ( - Self { - writer_id, - affinity_tx: Some(tx), - ..Self::default() - }, - rx, - ) - } - - fn affinity_sink( - &self, - partition: &RoutingPartitionId, - ) -> Option<(u64, Arc)> { - self.affinity_tx.clone().map(|tx| { - let sink = AffinityMeshSink { - partition: partition.clone(), - tx, - }; - ( - self.writer_id, - Arc::new(sink) as Arc, - ) - }) - } -} - #[derive(Debug, Clone)] pub struct SelectionServiceConfig { pub port: u16, @@ -326,6 +290,8 @@ pub struct SelectionCore { selection_cache: SelectionCache, tracking_hash: Arc, session_affinity: Option, + /// Writer id and publisher for session bindings the host carries to its peers. + session_binding_tx: Option<(u64, mpsc::Sender)>, /// Worker ids whose upsert fails with `Internal` before any catalog /// mutation, so membership tests can exercise per-worker error paths. #[cfg(test)] @@ -439,6 +405,7 @@ impl SelectionCore { selection_cache: SelectionCache::new(&cache_config), tracking_hash, session_affinity, + session_binding_tx: None, #[cfg(test)] fail_upsert_for: parking_lot::Mutex::default(), #[cfg(test)] @@ -448,6 +415,32 @@ impl SelectionCore { } } + /// Export session bindings through a bounded channel for the host to carry. + pub(super) fn with_session_binding_publisher( + mut self, + writer_id: u64, + ) -> (Self, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(AFFINITY_EVENT_CHANNEL_CAPACITY); + self.session_binding_tx = Some((writer_id, tx)); + (self, rx) + } + + fn session_binding_sink( + &self, + partition: &RoutingPartitionId, + ) -> Option<(u64, Arc)> { + self.session_binding_tx.as_ref().map(|(writer_id, tx)| { + let sink = AffinityMeshSink { + partition: partition.clone(), + tx: tx.clone(), + }; + ( + *writer_id, + Arc::new(sink) as Arc, + ) + }) + } + /// Cancel core-scoped tasks (KV-event listeners, scheduling, replica sync, /// periodic expiry) without cancelling the parent token. In-flight and /// queued selections then fail fast. diff --git a/lib/kv-router/src/services/selection/core/reservations.rs b/lib/kv-router/src/services/selection/core/reservations.rs index f12c59ad9e33..4914d197381d 100644 --- a/lib/kv-router/src/services/selection/core/reservations.rs +++ b/lib/kv-router/src/services/selection/core/reservations.rs @@ -387,7 +387,7 @@ impl SelectionCore { // initializer that may itself be waiting for this worker's capacity. let session = match (session_id.as_deref(), entry.affinity.get()) { (Some(session_id), Some(table)) => self - .hold_session(table, session_id, &key, None) + .hold_session(table, session_id, &key) .await? .map(|hold| (table, session_id, hold)), _ => None, diff --git a/lib/kv-router/src/services/selection/core/run.rs b/lib/kv-router/src/services/selection/core/run.rs index 064bd0764c10..27aceaa7cee9 100644 --- a/lib/kv-router/src/services/selection/core/run.rs +++ b/lib/kv-router/src/services/selection/core/run.rs @@ -42,8 +42,6 @@ //! (`dropped_selection_future_frees_its_booking`, //! `dropped_book_selection_during_routing_record_frees_booking_and_claim`). -use std::collections::HashSet; - use super::super::affinity::{SessionAffinityMode, validate_dispatch_target}; use super::hint::{hint_capable_partition, transfer_hint_for_selection}; @@ -308,9 +306,7 @@ impl SelectionCore { }; let affinity_target = match (&session, table) { (SessionBinding::Managed { session_id }, Some(table)) => { - affinity_hold = self - .hold_session(table, session_id, &key, allowed_worker_ids.as_ref()) - .await?; + affinity_hold = self.hold_session(table, session_id, &key).await?; affinity_hold.as_ref().and_then(Hold::target) } (SessionBinding::Query { session_id }, Some(table)) => table @@ -570,28 +566,24 @@ impl SelectionCore { } /// Hold `session_id` for a booking, re-initializing a binding whose target - /// this request can no longer schedule. `None` when the table is full: a + /// this partition can no longer schedule. `None` when the table is full: a /// router-side limit, not a client fault, so the request routes unpinned. pub(super) async fn hold_session( &self, table: &SessionAffinity, session_id: &str, key: &RoutingPartitionId, - allowed_worker_ids: Option<&HashSet>, ) -> Result, SelectionError> { let mut invalidations = 0; loop { match table.acquire(session_id, None).await { Ok(Hold::Bound { target, mut lease }) - if !self.catalog.is_schedulable(target, key) - || (table.mode() == SessionAffinityMode::Soft - && allowed_worker_ids - .is_some_and(|ids| !ids.contains(&target.worker_id))) => + if !self.catalog.is_schedulable(target, key) => { tracing::debug!( session_id, worker_id = target.worker_id, - "Session affinity target is unavailable or excluded; re-initializing" + "Session affinity target is not schedulable; re-initializing" ); lease.invalidate(); drop(lease); diff --git a/lib/kv-router/src/services/selection/core/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index 38828b398538..77e8d09fe9a5 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -2880,7 +2880,7 @@ async fn joined_failover_preserves_the_successful_bookings_binding(#[case] commi let mut context = std::task::Context::from_waker(std::task::Waker::noop()); let mut r2_request = session_reservation("r2", "s"); - r2_request.allowed_worker_ids = Some(HashSet::from([first.worker_id, replacement])); + r2_request.allowed_worker_ids = Some(HashSet::from([replacement])); let mut r2 = Box::pin(core.select_and_reserve(r2_request)); assert!(r2.as_mut().poll(&mut context).is_pending()); core.catalog @@ -2925,19 +2925,11 @@ async fn late_failover_mismatch_preserves_the_committed_replacement() { let key = default_key(); let entry = core.entry(&key).expect("entry"); let table = entry.affinity.get().expect("affinity table"); - let old_hold = core - .hold_session(table, "s", &key, None) - .await - .unwrap() - .unwrap(); + let old_hold = core.hold_session(table, "s", &key).await.unwrap().unwrap(); core.catalog .set_lifecycle(first.worker_id, WorkerLifecycle::Draining, Vec::new()); let replacement = if first.worker_id == 1 { 2 } else { 1 }; - let new_hold = core - .hold_session(table, "s", &key, None) - .await - .unwrap() - .unwrap(); + let new_hold = core.hold_session(table, "s", &key).await.unwrap().unwrap(); let lease = core .commit_session( table, @@ -3134,7 +3126,7 @@ async fn soft_mode_follows_the_dispatch() { } #[tokio::test] -async fn soft_binding_follows_policy_override() { +async fn soft_mode_follows_custom_policy() { // Bind to an eligible worker different from the custom policy's choice. let (policy, _) = capturing_policy_factory(); let core = core_with( @@ -3964,21 +3956,17 @@ async fn overflowing_worker_dp_range_patch_preserves_live_worker() { fn core_with_binding_stream( writer_id: u64, ) -> (SelectionCore, mpsc::Receiver) { - let (replication, rx) = HostReplication::affinity_channel(writer_id); - let core = core_with( + core_with( test_config(false), - SelectionHost { - replication, - ..SelectionHost::default() - }, + SelectionHost::default(), None, WorkerType::Aggregated, Some( SessionAffinityConfig::new(Duration::from_secs(10)) .with_mode(SessionAffinityMode::Soft), ), - ); - (core, rx) + ) + .with_session_binding_publisher(writer_id) } #[tokio::test(start_paused = true)] diff --git a/lib/kv-router/src/services/selection/core/workers.rs b/lib/kv-router/src/services/selection/core/workers.rs index 9326bfc648aa..3c168ed6d225 100644 --- a/lib/kv-router/src/services/selection/core/workers.rs +++ b/lib/kv-router/src/services/selection/core/workers.rs @@ -374,7 +374,7 @@ impl SelectionCore { replica_inbox, affinity: OnceCell::new(), replica_config: self.replica_config.clone(), - affinity_replication: self.host.replication.affinity_sink(&key), + affinity_replication: self.session_binding_sink(&key), })) })? .clone(); diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 1ead5ba38494..b64fa2a9dfaf 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -19,9 +19,7 @@ use crate::services::indexer::backend::IndexerPolicy; use crate::tracking_hash::TrackingHashContext; use super::affinity::{SessionAffinityConfig, SessionAffinityMode}; -use super::core::{ - HostReplication, KvIndexSource, SelectionCore, SelectionHost, SelectionServiceConfig, -}; +use super::core::{KvIndexSource, SelectionCore, SelectionHost, SelectionServiceConfig}; use super::error::SelectionError; use super::pending::SelectionCacheConfig; use super::types::{ @@ -168,26 +166,16 @@ impl SelectionServiceBuilder { } pub async fn build(mut self) -> anyhow::Result { - let session_binding_rx = if let Some(writer_id) = self.session_replication_writer_id { + if self.session_replication_writer_id.is_some() { anyhow::ensure!( self.session_affinity_ttl.is_some(), "session binding replication requires session affinity" ); anyhow::ensure!( - self.host.replication.affinity_tx.is_none(), - "session binding publisher already configured" + self.replica_sync_port.is_none(), + "host session binding replication cannot be combined with replica_sync" ); - let (replication, rx) = HostReplication::affinity_channel(writer_id); - self.host.replication.writer_id = replication.writer_id; - self.host.replication.affinity_tx = replication.affinity_tx; - Some(Mutex::new(Some(rx))) - } else { - None - }; - anyhow::ensure!( - self.replica_sync_port.is_none() || self.host.replication.affinity_tx.is_none(), - "host session binding replication cannot be combined with replica_sync" - ); + } self.kv_router_config .apply_policy_config() .map_err(anyhow::Error::msg)?; @@ -227,7 +215,7 @@ impl SelectionServiceBuilder { ); } let replica_config = replica_runtime.as_ref().map(ReplicaSyncRuntime::config); - let core = Arc::new(SelectionCore::new_inner( + let core = SelectionCore::new_inner( self.kv_router_config, self.indexer_threads, cancel_token.clone(), @@ -241,7 +229,15 @@ impl SelectionServiceBuilder { indexer_policy, self.session_affinity_ttl .map(|ttl| SessionAffinityConfig::new(ttl).with_mode(self.session_affinity_mode)), - )); + ); + let (core, session_binding_rx) = match self.session_replication_writer_id { + Some(writer_id) => { + let (core, rx) = core.with_session_binding_publisher(writer_id); + (core, Some(Mutex::new(Some(rx)))) + } + None => (core, None), + }; + let core = Arc::new(core); if recover_from_peers { match core.recover_indexer_from_peers(&self.indexer_peers).await { diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 5ecab3e11cff..5885b258e6d7 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -285,7 +285,6 @@ impl EmbeddedSelection { channels: replica_sync, request_leases, replica_worker_policy: ReplicaWorkerPolicy::LazyRegister, - ..HostReplication::default() }, }) .build() From e8f3b8825261767fadb54acd952d3c7a81a099de Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Sun, 25 Oct 2026 00:54:01 -0700 Subject: [PATCH 4/5] fix(bindings): pass new SelectionService args in lifecycle test Signed-off-by: Jeffrey Wang --- lib/bindings/python/rust/llm/kv.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index f20e969c88cd..49fe3d9a413d 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -1047,9 +1047,10 @@ mod selection_service_lifecycle_tests { #[test] fn idempotent_shutdown() { pyo3::prepare_freethreaded_python(); - let service = - Python::with_gil(|py| SelectionService::new(py, 1, None, None, None, None, None)) - .unwrap(); + let service = Python::with_gil(|py| { + SelectionService::new(py, 1, None, None, None, None, None, "hard", None) + }) + .unwrap(); Python::with_gil(|py| { service.shutdown(py); service.shutdown(py); From fcafcc74b11489e68e98da288a610ca03c518224 Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Mon, 26 Oct 2026 21:48:15 -0700 Subject: [PATCH 5/5] fix(kv-router): batch session binding stream, warn on drops, sync apply Co-Authored-By: Claude Opus 5.5 Signed-off-by: Jeffrey Wang --- lib/bindings/python/rust/llm/kv.rs | 21 +++------ lib/bindings/python/src/dynamo/_core.pyi | 5 +- lib/bindings/python/tests/test_kv_bindings.py | 47 +++++++++++++++---- .../src/services/common/replica_sync.rs | 42 +++++++++++++---- .../src/services/selection/core/mod.rs | 5 +- .../src/services/selection/service.rs | 6 ++- 6 files changed, 87 insertions(+), 39 deletions(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index b4f9ff038833..528b2ebe5218 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -702,10 +702,10 @@ impl SessionBindingStream { fn __anext__<'p>(&self, py: Python<'p>) -> PyResult> { let updates = self.updates.clone(); crate::future_into_py(py, async move { - let event = updates.recv().await.ok_or_else(|| { + let events = updates.recv_many().await.ok_or_else(|| { pyo3::exceptions::PyStopAsyncIteration::new_err("selection service shut down") })?; - Python::with_gil(|py| pythonize(py, &event).map(|o| o.unbind()).map_err(to_pyerr)) + Python::with_gil(|py| pythonize(py, &events).map(|o| o.unbind()).map_err(to_pyerr)) }) } } @@ -791,19 +791,12 @@ impl SelectionService { }) } - fn apply_session_binding<'p>( - &self, - py: Python<'p>, - update: PyObject, - ) -> PyResult> { - let event: selection::AffinityBindingEvent = + fn apply_session_binding(&self, py: Python<'_>, update: PyObject) -> PyResult<()> { + let event = depythonize(update.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?; - let service = self.inner.clone(); - crate::future_into_py(py, async move { - service - .apply_session_binding(event) - .map_err(selection_to_pyerr) - }) + self.inner + .apply_session_binding(event) + .map_err(selection_to_pyerr) } /// Await service shutdown without blocking the Python event loop. diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index b5af30cd1d1c..21290c85b76f 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -811,15 +811,16 @@ class SelectionService: """Await service shutdown without blocking the Python event loop.""" ... - def session_binding_updates(self) -> AsyncIterator[JsonLike]: + def session_binding_updates(self) -> AsyncIterator[List[JsonLike]]: """Stream session bindings this service commits or refreshes, for the host to broadcast to its peers. + Each item is every binding queued since the last one (at least one). Requires session_replication_writer_id. Ends on shutdown. """ ... - async def apply_session_binding(self, update: JsonLike) -> None: + def apply_session_binding(self, update: JsonLike) -> None: """Merge a versioned peer binding without echoing it. A binding this replica cannot place yet is dropped; the owner's next diff --git a/lib/bindings/python/tests/test_kv_bindings.py b/lib/bindings/python/tests/test_kv_bindings.py index 59ad01479fc2..8b1949ea1ad8 100644 --- a/lib/bindings/python/tests/test_kv_bindings.py +++ b/lib/bindings/python/tests/test_kv_bindings.py @@ -14,6 +14,7 @@ # limitations under the License. +import inspect import json import subprocess import sys @@ -347,11 +348,11 @@ async def test_selection_service_session_binding_stream(session_binding_services } ) assert "session_binding" not in selected - binding = json.loads(json.dumps(await anext(stream))) + (binding,) = json.loads(json.dumps(await anext(stream))) assert binding["writer_id"] == 11 - await peer.apply_session_binding(binding) + peer.apply_session_binding(binding) await source.free_reservation("origin") - assert await anext(stream) == binding # A repeated version refreshes TTL. + assert await anext(stream) == [binding] # A repeated version refreshes TTL. selected = await peer.select_and_reserve( { **partition, @@ -361,9 +362,9 @@ async def test_selection_service_session_binding_stream(session_binding_services } ) assert selected["worker_id"] == 1 - assert ( - await anext(peer_stream) == binding - ) # Local commits preserve imported versions. + assert await anext(peer_stream) == [ + binding + ] # Local commits preserve imported versions. # The default hard mode rejects an excluded binding; soft mode rebinds. request = { **partition, @@ -375,12 +376,38 @@ async def test_selection_service_session_binding_stream(session_binding_services with pytest.raises(ValueError, match="session s is bound to worker 1, not 2"): await source.select_and_reserve(request) assert (await peer.select_and_reserve(request))["worker_id"] == 2 - assert (await anext(peer_stream))["worker_id"] == 2 + assert (await anext(peer_stream))[0]["worker_id"] == 2 await source.shutdown_async() with pytest.raises(StopAsyncIteration): await anext(stream) with pytest.raises(SelectionServiceError): - await source.apply_session_binding(binding) + source.apply_session_binding(binding) + + +@pytest.mark.skipif( + SelectionService is None, + reason="SelectionService requires the select-service Cargo feature", +) +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_selection_service_session_binding_stream_drains_queue( + session_binding_services, +): + source, _ = session_binding_services + stream = source.session_binding_updates() + for session_id in ("a", "b", "c"): + await source.select_and_reserve( + { + "model_name": "model", + "routing_group": "host", + "token_ids": [1, 2, 3, 4], + "selection_id": session_id, + "session_id": session_id, + } + ) + # Bindings queued while the consumer was away arrive as one item. + batch = await anext(stream) + assert [binding["session_id"] for binding in batch] == ["a", "b", "c"] @pytest.mark.skipif( @@ -409,7 +436,9 @@ async def test_selection_service_malformed_payload_raises_value_error( try: # Malformed input surfaces as a plain ValueError, not a SelectionServiceError. with pytest.raises(ValueError): - await getattr(service, method)(payload) + result = getattr(service, method)(payload) + if inspect.isawaitable(result): + await result # Prompt booleans are not coerced by truthiness. with pytest.raises(ValueError, match="boolean"): await service.select({"token_ids": [1, 2, 3, 4], "is_eagle": "false"}) diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index ffe819a252d3..513a0fb3815e 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -65,8 +65,23 @@ impl AffinityBindingEvent { /// Publishes bindings from a [`super::super::selection::affinity::SessionAffinity`] /// into a transport channel. Best effort: a full channel drops the update. pub(crate) struct AffinityMeshSink { - pub(crate) partition: RoutingPartitionId, - pub(crate) tx: mpsc::Sender, + partition: RoutingPartitionId, + tx: mpsc::Sender, + full_drops: parking_lot::Mutex, +} + +#[cfg(feature = "standalone-selection")] +impl AffinityMeshSink { + pub(crate) fn new( + partition: RoutingPartitionId, + tx: mpsc::Sender, + ) -> Self { + Self { + partition, + tx, + full_drops: parking_lot::Mutex::default(), + } + } } #[cfg(feature = "standalone-selection")] @@ -80,8 +95,22 @@ impl AffinityReplicaSink for AffinityMeshSink { sequence: version.sequence, writer_id: version.writer_id, }; - if let Err(error) = self.tx.try_send(update) { - tracing::trace!(%error, "dropping best-effort session affinity replica update"); + match self.tx.try_send(update) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(_)) => { + let Some(dropped_bindings) = self.full_drops.lock().record(Instant::now()) else { + return; + }; + tracing::warn!( + partition = %self.partition, + capacity = self.tx.max_capacity(), + dropped_bindings_since_last_log = dropped_bindings, + "Session affinity replication channel full; dropping binding" + ); + } + Err(mpsc::error::TrySendError::Closed(_)) => { + tracing::trace!("dropping session affinity replica update; channel closed"); + } } } } @@ -183,10 +212,7 @@ impl ReplicaSyncConfig { partition: &RoutingPartitionId, ) -> Option> { self.affinity_tx.clone().map(|tx| { - Arc::new(AffinityMeshSink { - tx, - partition: partition.clone(), - }) as Arc + Arc::new(AffinityMeshSink::new(partition.clone(), tx)) as Arc }) } diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index f7863a0beee9..7533037020e8 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -430,10 +430,7 @@ impl SelectionCore { partition: &RoutingPartitionId, ) -> Option<(u64, Arc)> { self.session_binding_tx.as_ref().map(|(writer_id, tx)| { - let sink = AffinityMeshSink { - partition: partition.clone(), - tx: tx.clone(), - }; + let sink = AffinityMeshSink::new(partition.clone(), tx.clone()); ( *writer_id, Arc::new(sink) as Arc, diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index b64fa2a9dfaf..a18f45c01b50 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -356,12 +356,14 @@ pub struct SessionBindingUpdates { } impl SessionBindingUpdates { - pub async fn recv(&self) -> Option { + pub async fn recv_many(&self) -> Option> { let mut rx = self.rx.lock().await; + let limit = rx.max_capacity(); + let mut events = Vec::new(); tokio::select! { biased; _ = self.cancel_token.cancelled() => None, - event = rx.recv() => event, + received = rx.recv_many(&mut events, limit) => (received > 0).then_some(events), } } }