diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 9a2f8def87d3..9c15b335010d 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -396,6 +396,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 4fea01a3b6ba..528b2ebe5218 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -686,12 +686,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 events = updates.recv_many().await.ok_or_else(|| { + pyo3::exceptions::PyStopAsyncIteration::new_err("selection service shut down") + })?; + Python::with_gil(|py| pythonize(py, &events).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<'_>, @@ -701,7 +725,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( @@ -719,7 +748,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); } @@ -728,6 +758,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)?; @@ -746,6 +779,26 @@ 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(&self, py: Python<'_>, update: PyObject) -> PyResult<()> { + let event = + depythonize(update.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?; + self.inner + .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); @@ -1020,9 +1073,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); diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index ed0e899f227e..21290c85b76f 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -783,11 +783,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. """ ... @@ -801,6 +807,27 @@ class SelectionService: """ ... + async def shutdown_async(self) -> None: + """Await service shutdown without blocking the Python event loop.""" + ... + + 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. + """ + ... + + 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.""" ... diff --git a/lib/bindings/python/tests/test_kv_bindings.py b/lib/bindings/python/tests/test_kv_bindings.py index 9d904a2f98f6..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 @@ -291,20 +292,153 @@ 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, + 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: + 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.timeout(5) @pytest.mark.asyncio -async def test_selection_service_malformed_payload_raises_value_error(monkeypatch): +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"} + selected = await source.select_and_reserve( + { + **partition, + "token_ids": [1, 2, 3, 4], + "selection_id": "origin", + "session_id": "s", + "allowed_worker_ids": [1], + } + ) + assert "session_binding" not in selected + (binding,) = json.loads(json.dumps(await anext(stream))) + assert binding["writer_id"] == 11 + peer.apply_session_binding(binding) + await source.free_reservation("origin") + assert await anext(stream) == [binding] # A repeated version refreshes TTL. + selected = await peer.select_and_reserve( + { + **partition, + "token_ids": [5, 6, 7, 8], + "selection_id": "next", + "session_id": "s", + } + ) + assert selected["worker_id"] == 1 + 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))[0]["worker_id"] == 2 + await source.shutdown_async() + with pytest.raises(StopAsyncIteration): + await anext(stream) + with pytest.raises(SelectionServiceError): + 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( + 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( + "apply_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"}) + 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 0fef46e5b038..513a0fb3815e 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -33,11 +33,11 @@ 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)] -pub(crate) struct AffinityBindingEvent { +pub struct AffinityBindingEvent { #[serde(flatten)] pub partition: RoutingPartitionId, pub session_id: String, @@ -63,10 +63,25 @@ 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 { +/// into a transport channel. Best effort: a full channel drops the update. +pub(crate) struct AffinityMeshSink { 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 1ff6688d438f..7533037020e8 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) })?; @@ -286,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)] @@ -399,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)] @@ -408,6 +415,29 @@ 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::new(partition.clone(), 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. @@ -418,7 +448,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/tests.rs b/lib/kv-router/src/services/selection/core/tests.rs index d0f52b54ecc4..77e8d09fe9a5 100644 --- a/lib/kv-router/src/services/selection/core/tests.rs +++ b/lib/kv-router/src/services/selection/core/tests.rs @@ -3125,6 +3125,42 @@ async fn soft_mode_follows_the_dispatch() { assert_eq!(bound_worker(&core, "s"), Some(other)); } +#[tokio::test] +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( + test_config(false), + SelectionHost::default(), + Some(policy), + WorkerType::Aggregated, + Some( + SessionAffinityConfig::new(Duration::from_secs(10)) + .with_mode(SessionAffinityMode::Soft), + ), + ); + 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)); + let next = core + .select_and_reserve(session_reservation("next-turn", "s")) + .await + .unwrap(); + assert_eq!(next.worker_id, preferred); + assert_eq!(bound_worker(&core, "s"), Some(preferred)); +} + #[tokio::test] async fn replicated_binding_steers_a_new_session_and_frees_with_the_booking() { let core = core_with_session_affinity(); @@ -3916,3 +3952,103 @@ 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) { + core_with( + test_config(false), + SelectionHost::default(), + None, + WorkerType::Aggregated, + Some( + SessionAffinityConfig::new(Duration::from_secs(10)) + .with_mode(SessionAffinityMode::Soft), + ), + ) + .with_session_binding_publisher(writer_id) +} + +#[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..3c168ed6d225 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.session_binding_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/service.rs b/lib/kv-router/src/services/selection/service.rs index 057b328dcbf5..a18f45c01b50 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,12 +13,12 @@ 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::affinity::{SessionAffinityConfig, SessionAffinityMode}; use super::core::{KvIndexSource, SelectionCore, SelectionHost, SelectionServiceConfig}; use super::error::SelectionError; use super::pending::SelectionCacheConfig; @@ -40,6 +42,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 +85,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 +131,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 +166,16 @@ impl SelectionServiceBuilder { } pub async fn build(mut self) -> anyhow::Result { + 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.replica_sync_port.is_none(), + "host session binding replication cannot be combined with replica_sync" + ); + } self.kv_router_config .apply_policy_config() .map_err(anyhow::Error::msg)?; @@ -185,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(), @@ -197,8 +227,17 @@ 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)), + ); + 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 { @@ -241,6 +280,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 +341,66 @@ 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_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, + received = rx.recv_many(&mut events, limit) => (received > 0).then_some(events), + } + } +} + 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 +408,7 @@ impl SelectionService { ) -> Self { let cancel_token = CancellationToken::new(); Self { + session_binding_rx: None, core: Arc::new( SelectionCore::try_new_local( kv_router_config, diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index a30e40d514ae..da221128b12b 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 {