diff --git a/lib/kv-router/src/services/common/replica_sync.rs b/lib/kv-router/src/services/common/replica_sync.rs index 810d0ee06885..752112d9df07 100644 --- a/lib/kv-router/src/services/common/replica_sync.rs +++ b/lib/kv-router/src/services/common/replica_sync.rs @@ -24,67 +24,24 @@ use crate::services::common::zmq::{ create_bound_pub_socket, create_sub_socket_topics, validate_endpoint, }; #[cfg(feature = "standalone-selection")] -use crate::services::selection::affinity::{AffinityReplicaSink, AffinityTarget, AffinityVersion}; +pub(crate) use crate::services::selection::affinity::AffinityBindingEvent; +#[cfg(feature = "standalone-selection")] +use crate::services::selection::affinity::{AffinityEventSink, AffinityReplicaSink}; pub(crate) const REPLICA_EVENT_CHANNEL_CAPACITY: usize = 100_000; const PEER_COMMAND_CHANNEL_CAPACITY: usize = 64; 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"; +/// Session-affinity bindings ride the same mesh on their own topics, so peers +/// that do not subscribe to them never see them. The payload is the shared +/// [`AffinityBindingEvent`] on both topics; v2 is the topic every host's +/// transport uses, v1 is published and applied alongside it for one release +/// so a mixed-version mesh stays converged. +// Compatibility with v1.6 selection replicas during v1.7 rolling upgrades. +// TODO(v1.8): Publish and subscribe on `AFFINITY_TOPIC_V2` only. +const AFFINITY_TOPIC_V1: &[u8] = b"dynamo.session-affinity.v1"; +const AFFINITY_TOPIC_V2: &[u8] = b"dynamo.session-affinity.v2"; const AFFINITY_EVENT_CHANNEL_CAPACITY: usize = 4_096; -/// One replicated session binding. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub(crate) struct AffinityBindingEvent { - #[serde(flatten)] - pub partition: RoutingPartitionId, - pub session_id: String, - pub worker_id: u64, - pub dp_rank: Option, - pub sequence: u64, - pub writer_id: u64, -} - -#[cfg(feature = "standalone-selection")] -impl AffinityBindingEvent { - pub(crate) fn target(&self) -> AffinityTarget { - AffinityTarget::new(self.worker_id, self.dp_rank) - } - - pub(crate) fn version(&self) -> AffinityVersion { - AffinityVersion { - sequence: self.sequence, - writer_id: self.writer_id, - } - } -} - -#[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, -} - -#[cfg(feature = "standalone-selection")] -impl AffinityReplicaSink for AffinityMeshSink { - fn publish(&self, session_id: &str, target: AffinityTarget, version: AffinityVersion) { - let update = AffinityBindingEvent { - partition: self.partition.clone(), - session_id: session_id.to_string(), - worker_id: target.worker_id, - dp_rank: target.dp_rank, - 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"); - } - } -} - #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct ScopedReplicaEvent { pub model_name: String, @@ -176,16 +133,17 @@ impl ReplicaSyncConfig { #[cfg(feature = "standalone-selection")] /// Sink that publishes session bindings into the mesh, when this runtime - /// carries them. + /// carries them. Best effort: a full channel drops the update. pub(crate) fn affinity_sink( &self, partition: &RoutingPartitionId, ) -> Option> { self.affinity_tx.clone().map(|tx| { - Arc::new(AffinityMeshSink { - tx, - partition: partition.clone(), - }) as Arc + Arc::new(AffinityEventSink::new(partition.clone(), move |event| { + if let Err(error) = tx.try_send(event) { + tracing::trace!(%error, "dropping best-effort session affinity replica update"); + } + })) as Arc }) } @@ -550,8 +508,10 @@ pub(crate) fn start_replica_publisher( continue; } }; - if let Err(error) = socket.send_multipart(frames).await { - tracing::error!(event = %message, error = %error, "Failed to publish replica message"); + for frames in frames { + if let Err(error) = socket.send_multipart(frames).await { + tracing::error!(event = %message, error = %error, "Failed to publish replica message"); + } } } }); @@ -565,12 +525,24 @@ enum OutboundReplicaMessage { } impl OutboundReplicaMessage { - fn encode(&self) -> Result>, rmp_serde::encode::Error> { - let (topic, payload) = match self { - Self::Sequence(event) => (REPLICA_TOPIC, rmp_serde::to_vec_named(event)?), - Self::Affinity(binding) => (AFFINITY_TOPIC, rmp_serde::to_vec_named(binding)?), - }; - Ok(vec![topic.to_vec(), payload]) + /// One multipart message per topic the event goes out on; a binding goes + /// out on both affinity topics for the compatibility window. + fn encode(&self) -> Result>>, rmp_serde::encode::Error> { + Ok(match self { + Self::Sequence(event) => { + vec![vec![ + REPLICA_TOPIC.to_vec(), + rmp_serde::to_vec_named(event)?, + ]] + } + Self::Affinity(binding) => { + let payload = rmp_serde::to_vec_named(binding)?; + vec![ + vec![AFFINITY_TOPIC_V2.to_vec(), payload.clone()], + vec![AFFINITY_TOPIC_V1.to_vec(), payload], + ] + } + }) } } @@ -707,7 +679,7 @@ impl PeerManager { A: Fn(AffinityBindingEvent) + Send + Sync + 'static, { let topics: &[&[u8]] = if handle_affinity.is_some() { - &[REPLICA_TOPIC, AFFINITY_TOPIC] + &[REPLICA_TOPIC, AFFINITY_TOPIC_V1, AFFINITY_TOPIC_V2] } else { &[REPLICA_TOPIC] }; @@ -871,7 +843,9 @@ fn handle_replica_message( tracing::debug!("Dropping malformed active-sequence replica payload: {error}"); } }, - AFFINITY_TOPIC => match ( + // Both topics carry the same payload; applying a binding twice is + // idempotent (the second apply refreshes the same version). + AFFINITY_TOPIC_V1 | AFFINITY_TOPIC_V2 => match ( handle_affinity, rmp_serde::from_slice::(payload), ) { diff --git a/lib/kv-router/src/services/selection/affinity/README.md b/lib/kv-router/src/services/selection/affinity/README.md index bce2203c5310..8014d7fb373f 100644 --- a/lib/kv-router/src/services/selection/affinity/README.md +++ b/lib/kv-router/src/services/selection/affinity/README.md @@ -39,6 +39,14 @@ Every host resolves through `AffinityResolver`. The frontend keeps `HostAffinity Where the hosts differ on a live `Hard` mismatch (the bound worker is live but the request's own constraints steered selection elsewhere): the core rejects at commit and the table drops the binding, so the session re-binds on its next request; the frontend rejects before dispatch with `check_dispatch` and releases its hold without touching the binding, since nothing ran and the request, not the session, caused the mismatch. +## Replication + +One wire schema, `AffinityBindingEvent` (`replication.rs`): the partition (flattened as `model_name` and `routing_group`), the session id, the target, and the version (`sequence`, `writer_id`). Every table publishes through `AffinityEventSink`, which builds the event for its partition and hands it to the host's transport; every transport feeds received events to `AffinityResolver::apply_replica_event`, the one applier. It ignores the replica's own writer id, ignores other partitions, advances the replica clock, and applies the binding only if the host can schedule its worker (`TargetLiveness`). + +One writer-id rule: the id installed with the table's replication sink. Frontends use their discovery instance id (stable across restarts of the same instance); the standalone service uses a random non-zero process id, since it has no discovery. + +Two transports stay: the selection service's ZMQ peer mesh (`services/common/replica_sync.rs`, topics `dynamo.session-affinity.v1` and `.v2`) and the frontend's runtime event plane (`lib/llm/src/session_affinity/replica_sync.rs`, subjects `session_affinity_events` for the old partition-less payload and `session_affinity_events_v2` for the shared schema). For one release both transports publish and apply both versions, so mixed-version replicas converge; applying the same binding twice is idempotent (the second apply refreshes the same version). The v1 forms carry a removal TODO. + ## Errors Hosts map `AffinityError` onto their own error types. The core maps `InvalidArgument` (the request contradicts the binding or exceeds the session-id limit) to `BadRequest`, `ResourceExhausted` to `NotReady`, `Cancelled` to scheduler shutdown, and `Dropped` to `Internal`. A full table never fails `resolve`: the request routes without affinity and `full_table_fallbacks` counts it (a plain counter on the resolver; not exported as a metric). diff --git a/lib/kv-router/src/services/selection/affinity/mod.rs b/lib/kv-router/src/services/selection/affinity/mod.rs index c03972e33f1d..59e77dd42151 100644 --- a/lib/kv-router/src/services/selection/affinity/mod.rs +++ b/lib/kv-router/src/services/selection/affinity/mod.rs @@ -6,8 +6,10 @@ //! [`AffinityResolver`] drives it for one request; `README.md` in this //! directory maps the routing hosts onto both. +mod replication; mod resolver; mod table; +pub use replication::{AffinityBindingEvent, AffinityEventSink, ReplicaEventDisposition}; pub use resolver::{AffinityResolver, Resolution, TargetLiveness}; pub use table::*; diff --git a/lib/kv-router/src/services/selection/affinity/replication.rs b/lib/kv-router/src/services/selection/affinity/replication.rs new file mode 100644 index 000000000000..af65baaaf89f --- /dev/null +++ b/lib/kv-router/src/services/selection/affinity/replication.rs @@ -0,0 +1,327 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Replicated session bindings: the one wire schema every host publishes and +//! applies, and the one applier. Transports are the hosts' (the selection +//! service's ZMQ peer mesh, the frontend's runtime event plane); each carries +//! [`AffinityBindingEvent`] and feeds received events to +//! [`AffinityResolver::apply_replica_event`]. + +use std::sync::Arc; + +use serde::{Deserialize, Serialize}; + +use super::resolver::AffinityResolver; +use super::table::{AffinityReplicaSink, AffinityTarget, AffinityVersion, ReplicaApplyOutcome}; +use crate::identity::RoutingPartitionId; + +/// One replicated session binding. The partition scopes the session id: a +/// binding for another partition is ignored on receipt. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct AffinityBindingEvent { + #[serde(flatten)] + pub partition: RoutingPartitionId, + pub session_id: String, + pub worker_id: u64, + pub dp_rank: Option, + pub sequence: u64, + /// The publishing replica's writer id: its discovery instance id where + /// the host has one (frontends), otherwise a random non-zero process id + /// (the standalone service). The applier ignores its own writer id. + pub writer_id: u64, +} + +impl AffinityBindingEvent { + pub fn new( + partition: &RoutingPartitionId, + session_id: &str, + target: AffinityTarget, + version: AffinityVersion, + ) -> Self { + Self { + partition: partition.clone(), + session_id: session_id.to_string(), + worker_id: target.worker_id, + dp_rank: target.dp_rank, + sequence: version.sequence, + writer_id: version.writer_id, + } + } + + pub fn target(&self) -> AffinityTarget { + AffinityTarget::new(self.worker_id, self.dp_rank) + } + + pub fn version(&self) -> AffinityVersion { + AffinityVersion { + sequence: self.sequence, + writer_id: self.writer_id, + } + } +} + +/// A table sink that turns each published binding into an +/// [`AffinityBindingEvent`] for `partition` and hands it to the transport. +/// Best effort: the transport drops what it cannot queue. +pub struct AffinityEventSink { + partition: RoutingPartitionId, + publish: Arc, +} + +impl AffinityEventSink { + pub fn new( + partition: RoutingPartitionId, + publish: impl Fn(AffinityBindingEvent) + Send + Sync + 'static, + ) -> Self { + Self { + partition, + publish: Arc::new(publish), + } + } +} + +impl AffinityReplicaSink for AffinityEventSink { + fn publish(&self, session_id: &str, target: AffinityTarget, version: AffinityVersion) { + (self.publish)(AffinityBindingEvent::new( + &self.partition, + session_id, + target, + version, + )); + } +} + +/// What the applier did with a received binding. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ReplicaEventDisposition { + Applied(ReplicaApplyOutcome), + /// Published by this replica. + OwnWriter, + /// For another partition than the receiving resolver's. + OtherPartition, + /// The host cannot schedule the bound worker; the version was still + /// observed so this replica's own versions stay ahead. + UnknownWorker, +} + +impl AffinityResolver { + /// Apply a binding another replica published, for the resolver serving + /// `partition`. Every host's transport ends here. + pub fn apply_replica_event( + &self, + partition: &RoutingPartitionId, + event: AffinityBindingEvent, + ) -> ReplicaEventDisposition { + if event.writer_id == self.table().writer_id() { + return ReplicaEventDisposition::OwnWriter; + } + if event.partition != *partition { + return ReplicaEventDisposition::OtherPartition; + } + self.table().observe_replica_sequence(event.sequence); + let (target, version) = (event.target(), event.version()); + if !self.is_schedulable(target) { + tracing::trace!( + key = %event.partition, + worker_id = event.worker_id, + "Dropping session affinity replica update: worker not schedulable here" + ); + return ReplicaEventDisposition::UnknownWorker; + } + let outcome = self + .table() + .apply_replica_update(event.session_id, target, version); + tracing::trace!( + worker_id = event.worker_id, + ?outcome, + "Applied session affinity replica update" + ); + ReplicaEventDisposition::Applied(outcome) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::sync::Mutex; + use std::time::Duration; + + use super::super::resolver::TargetLiveness; + use super::super::table::{SessionAffinity, SessionAffinityConfig}; + use super::*; + use crate::protocols::WorkerId; + + struct Live(HashSet); + + impl TargetLiveness for Live { + fn is_schedulable(&self, target: AffinityTarget) -> bool { + self.0.contains(&target.worker_id) + } + } + + fn partition_id(name: &str) -> RoutingPartitionId { + RoutingPartitionId::new(name, "default") + } + + fn resolver(live: &[WorkerId]) -> Arc { + let table = + SessionAffinity::with_config(SessionAffinityConfig::new(Duration::from_secs(60))) + .expect("table"); + Arc::new(AffinityResolver::new( + table, + Arc::new(Live(live.iter().copied().collect())), + )) + } + + /// Install a sink that records every event the table publishes. + fn record_events( + resolver: &AffinityResolver, + writer_id: u64, + partition: &RoutingPartitionId, + ) -> Arc>> { + let events = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&events); + assert!(resolver.table().enable_replication( + writer_id, + Arc::new(AffinityEventSink::new(partition.clone(), move |event| { + recorded.lock().unwrap().push(event); + })), + )); + events + } + + fn event( + partition: &RoutingPartitionId, + writer_id: u64, + sequence: u64, + ) -> AffinityBindingEvent { + AffinityBindingEvent { + partition: partition.clone(), + session_id: "s".to_string(), + worker_id: 1, + dp_rank: Some(0), + sequence, + writer_id, + } + } + + #[test] + fn schema_round_trips_with_the_partition_flattened() { + let event = event(&partition_id("model"), 7, 42); + let json = serde_json::to_value(&event).expect("encode"); + assert_eq!(json["model_name"], "model"); + assert_eq!(json["routing_group"], "default"); + assert_eq!(json["session_id"], "s"); + assert_eq!(json["writer_id"], 7); + let decoded: AffinityBindingEvent = serde_json::from_value(json).expect("decode"); + assert_eq!(decoded, event); + let bytes = rmp_serde::to_vec_named(&event).expect("msgpack encode"); + let decoded: AffinityBindingEvent = rmp_serde::from_slice(&bytes).expect("msgpack decode"); + assert_eq!(decoded, event); + } + + #[tokio::test] + async fn applier_scopes_by_partition_writer_and_liveness() { + let partition = partition_id("model"); + let resolver = resolver(&[1]); + let _events = record_events(&resolver, 7, &partition); + + assert_eq!( + resolver.apply_replica_event(&partition, event(&partition, 7, 1)), + ReplicaEventDisposition::OwnWriter + ); + assert_eq!( + resolver.apply_replica_event(&partition, event(&partition_id("other"), 9, 2)), + ReplicaEventDisposition::OtherPartition + ); + let mut unknown = event(&partition, 9, 3); + unknown.worker_id = 5; + assert_eq!( + resolver.apply_replica_event(&partition, unknown), + ReplicaEventDisposition::UnknownWorker + ); + assert!( + resolver.table().next_version().sequence > 3, + "an unapplied event still advances the replica clock" + ); + assert_eq!( + resolver.apply_replica_event(&partition, event(&partition, 9, 4)), + ReplicaEventDisposition::Applied(ReplicaApplyOutcome::Inserted) + ); + assert_eq!( + resolver.table().query_target("s", None).expect("query"), + Some(AffinityTarget::new(1, Some(0))) + ); + } + + /// Two replicas with different liveness sources (a frontend and a + /// selection service, in effect) converge on one binding through the + /// shared schema, in both directions. + #[tokio::test] + async fn replicas_converge_through_the_shared_schema() { + let partition = partition_id("model"); + let frontend = resolver(&[1, 2]); + let service = resolver(&[1, 2]); + let frontend_events = record_events(&frontend, 100, &partition); + let service_events = record_events(&service, 200, &partition); + + // The frontend binds; the service applies what it published. + let hold = frontend + .resolve("s", None, std::future::pending()) + .await + .expect("resolve") + .hold + .expect("hold"); + drop( + frontend + .commit(hold, AffinityTarget::new(1, Some(0))) + .expect("commit"), + ); + let published: Vec<_> = frontend_events.lock().unwrap().drain(..).collect(); + assert!(!published.is_empty()); + for event in published { + assert_eq!(event.partition, partition); + assert_eq!(event.writer_id, 100); + assert!(matches!( + service.apply_replica_event(&partition, event), + ReplicaEventDisposition::Applied(_) + )); + } + assert_eq!( + service.table().query_target("s", None).expect("query"), + Some(AffinityTarget::new(1, Some(0))) + ); + + // The service re-binds after the worker departs; the frontend follows. + let hold = service + .resolve("s", None, std::future::pending()) + .await + .expect("resolve") + .hold + .expect("hold"); + let Err(_) = service.commit(hold, AffinityTarget::new(2, Some(0))) else { + panic!("hard mismatch while the worker is live is rejected"); + }; + let hold = service + .resolve("s", None, std::future::pending()) + .await + .expect("resolve") + .hold + .expect("hold"); + drop( + service + .commit(hold, AffinityTarget::new(2, Some(0))) + .expect("rebind"), + ); + let published: Vec<_> = service_events.lock().unwrap().drain(..).collect(); + for event in published { + assert_eq!(event.writer_id, 200); + frontend.apply_replica_event(&partition, event); + } + assert_eq!( + frontend.table().query_target("s", None).expect("query"), + Some(AffinityTarget::new(2, Some(0))), + "the frontend converges on the service's newer binding" + ); + } +} diff --git a/lib/kv-router/src/services/selection/affinity/resolver.rs b/lib/kv-router/src/services/selection/affinity/resolver.rs index 5d4739519382..f64318ad02c7 100644 --- a/lib/kv-router/src/services/selection/affinity/resolver.rs +++ b/lib/kv-router/src/services/selection/affinity/resolver.rs @@ -76,6 +76,11 @@ impl AffinityResolver { &self.table } + /// Whether the host can schedule `target` right now. + pub fn is_schedulable(&self, target: AffinityTarget) -> bool { + self.liveness.is_schedulable(target) + } + /// Requests that ran without a session binding because the table was /// full: at `resolve`, or at a failover `commit` that found no room to /// re-bind. diff --git a/lib/kv-router/src/services/selection/affinity/table.rs b/lib/kv-router/src/services/selection/affinity/table.rs index b80f617c6263..96eb63bf2ebb 100644 --- a/lib/kv-router/src/services/selection/affinity/table.rs +++ b/lib/kv-router/src/services/selection/affinity/table.rs @@ -346,6 +346,12 @@ impl SessionAffinity { self.inner.mode } + /// This replica's writer id, installed with its replication sink; zero + /// until then. + pub fn writer_id(&self) -> u64 { + self.inner.writer_id.load(Ordering::Relaxed) + } + /// Acquire `session_id` for a request. `requested_target` is an explicit /// pin the request carries; a bound session must agree with it. pub fn try_acquire( diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index c40be1307dbc..36367f98c308 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -508,7 +508,9 @@ impl SelectionCore { } } - /// Apply a session binding a replica published. + /// Apply a session binding a replica published, through the partition's + /// resolver (which drops the publisher's own events and bindings to + /// workers this partition cannot schedule). pub(crate) fn dispatch_affinity_event(&self, event: AffinityBindingEvent) { let Some(entry) = self.entry(&event.partition) else { tracing::trace!( @@ -517,40 +519,14 @@ impl SelectionCore { ); return; }; - let Some(table) = entry.affinity.get().map(|resolver| resolver.table()) else { + let Some(resolver) = entry.affinity.get() else { tracing::trace!( key = %event.partition, "Dropping session affinity replica update: no affinity table" ); return; }; - if self - .replica_config - .as_ref() - .is_some_and(|config| config.process_id() == event.writer_id) - { - return; - } - table.observe_replica_sequence(event.sequence); - if self - .catalog - .get(event.worker_id) - .is_none_or(|record| record.key() != event.partition) - { - tracing::trace!( - key = %event.partition, - worker_id = event.worker_id, - "Dropping session affinity replica update: worker not in partition" - ); - return; - } - let (target, version, worker_id) = (event.target(), event.version(), event.worker_id); - let outcome = table.apply_replica_update(event.session_id, target, version); - tracing::trace!( - worker_id, - ?outcome, - "Applied session affinity replica update" - ); + resolver.apply_replica_event(&entry.key, event); } fn ready_entry(&self, key: &RoutingPartitionId) -> Result, SelectionError> { diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index 7b778e6bfcb3..6a3bd68c4e10 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -391,7 +391,12 @@ impl EmbeddedSelection { )?; Ok(self .affinity - .get_or_init(|| crate::session_affinity::HostAffinity::from_resolver(resolver)) + .get_or_init(|| { + crate::session_affinity::HostAffinity::from_resolver( + resolver, + self.partition.key().clone(), + ) + }) .clone()) } diff --git a/lib/llm/src/session_affinity/host.rs b/lib/llm/src/session_affinity/host.rs index 7af6f8286b04..65f7f15f4af7 100644 --- a/lib/llm/src/session_affinity/host.rs +++ b/lib/llm/src/session_affinity/host.rs @@ -14,6 +14,7 @@ use std::{ time::Duration, }; +use dynamo_kv_router::RoutingPartitionId; use dynamo_kv_router::services::selection::affinity::{ AffinityError, AffinityLease, AffinityResolver, SessionAffinity, SessionAffinityConfig, TargetLiveness, @@ -50,6 +51,9 @@ pub(crate) fn from_table(target: TableTarget) -> AffinityTarget { struct Inner { resolver: Arc, + /// The partition replicated bindings are scoped to: the embedded + /// partition's key for KV routing, the default key for builtin modes. + partition: RoutingPartitionId, replica: tokio::sync::OnceCell, } @@ -64,10 +68,14 @@ pub struct HostAffinity { impl HostAffinity { /// Wrap a resolver another component owns (an embedded selection partition). - pub(crate) fn from_resolver(resolver: Arc) -> Self { + pub(crate) fn from_resolver( + resolver: Arc, + partition: RoutingPartitionId, + ) -> Self { Self { inner: Arc::new(Inner { resolver, + partition, replica: tokio::sync::OnceCell::new(), }), } @@ -75,7 +83,8 @@ impl HostAffinity { /// A resolver over this host's own table, for hosts without a selection /// partition (builtin routing modes). `liveness` is the host's view of - /// which bound workers can still take requests. + /// which bound workers can still take requests; bindings replicate under + /// the default partition key. pub fn standalone( ttl: Duration, mode: SessionAffinityMode, @@ -83,9 +92,10 @@ impl HostAffinity { ) -> Result { let table = SessionAffinity::with_config(SessionAffinityConfig::new(ttl).with_mode(mode)) .map_err(affinity_error)?; - Ok(Self::from_resolver(Arc::new(AffinityResolver::new( - table, liveness, - )))) + Ok(Self::from_resolver( + Arc::new(AffinityResolver::new(table, liveness)), + crate::kv_router::embedded::embedded_partition_key(None), + )) } /// [`Self::standalone`] with runtime replication over `client`'s event @@ -104,13 +114,19 @@ impl HostAffinity { Ok(Some(affinity)) } + /// Replicate this host's bindings over `client`'s event plane and apply + /// the bindings other frontends publish for the same partition. pub(crate) async fn enable_replica_sync(&self, client: Client) -> Result<(), Error> { self.inner .replica .get_or_try_init(|| async { let table = self.inner.resolver.table(); - let (replica, router_id) = - ReplicaSyncRuntime::start(client, table.downgrade()).await?; + let (replica, router_id) = ReplicaSyncRuntime::start( + client, + Arc::downgrade(&self.inner.resolver), + self.inner.partition.clone(), + ) + .await?; if !table.enable_replication(router_id, replica.sink()) { return Err(anyhow::anyhow!( "session affinity table already has a replica sink installed" diff --git a/lib/llm/src/session_affinity/liveness.rs b/lib/llm/src/session_affinity/liveness.rs index 0a8e41921fec..55539d1e5b6f 100644 --- a/lib/llm/src/session_affinity/liveness.rs +++ b/lib/llm/src/session_affinity/liveness.rs @@ -32,6 +32,17 @@ impl DiscoveryLiveness { #[cfg(test)] pub(crate) struct AlwaysLive; +/// Exactly these workers are live; for tests of liveness-dependent paths. +#[cfg(test)] +pub(crate) struct LiveWorkers(pub(crate) std::collections::HashSet); + +#[cfg(test)] +impl TargetLiveness for LiveWorkers { + fn is_schedulable(&self, target: TableTarget) -> bool { + self.0.contains(&target.worker_id) + } +} + #[cfg(test)] impl TargetLiveness for AlwaysLive { fn is_schedulable(&self, _target: TableTarget) -> bool { diff --git a/lib/llm/src/session_affinity/mod.rs b/lib/llm/src/session_affinity/mod.rs index ce47d2906379..77b2ce18850a 100644 --- a/lib/llm/src/session_affinity/mod.rs +++ b/lib/llm/src/session_affinity/mod.rs @@ -18,9 +18,9 @@ pub use host::{AffinityTarget, HostAffinity, explicit_target}; pub(crate) use host::{ affinity_error, affinity_id, from_table, invalid_argument, to_table, tracked_stream, }; -#[cfg(test)] -pub(crate) use liveness::AlwaysLive; pub(crate) use liveness::DiscoveryLiveness; +#[cfg(test)] +pub(crate) use liveness::{AlwaysLive, LiveWorkers}; pub type LlmResponse = crate::types::Annotated; diff --git a/lib/llm/src/session_affinity/replica_sync.rs b/lib/llm/src/session_affinity/replica_sync.rs index 0ef8fc6d1be2..ed06c77c07d2 100644 --- a/lib/llm/src/session_affinity/replica_sync.rs +++ b/lib/llm/src/session_affinity/replica_sync.rs @@ -1,35 +1,48 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::sync::Arc; +//! Replication of session bindings between frontends over the runtime event +//! plane. The schema and the applier are the shared ones +//! (`AffinityBindingEvent`, `AffinityResolver::apply_replica_event`); this +//! module is the transport: publishers, subscribers, and the direct-ZMQ +//! fan-in when the deployment uses it. + +use std::sync::{Arc, Weak}; use anyhow::{Context, Result}; use dynamo_runtime::{ - component::{Client, Instance}, + component::Client, discovery::EventTransportKind, traits::DistributedRuntimeProvider, transports::event_plane::{ Codec, EventPublisher, EventSubscriber, ValidatedEnvelope, uses_direct_zmq, }, }; -use serde::{Deserialize, Serialize}; -use tokio::{ - sync::{mpsc, watch}, - task::JoinHandle, -}; +use serde::{Deserialize, Serialize, de::DeserializeOwned}; +use tokio::{sync::mpsc, task::JoinHandle}; use tokio_util::sync::CancellationToken; use crate::direct_zmq_fan_in::{ ContinuityMode, FanInEvent, FanInObservation, start_direct_zmq_fan_in, }; +use dynamo_kv_router::RoutingPartitionId; use dynamo_kv_router::services::selection::affinity::{ - AffinityReplicaSink, AffinityTarget, AffinityVersion, WeakSessionAffinity, + AffinityBindingEvent, AffinityEventSink, AffinityReplicaSink, AffinityResolver, }; +/// The v1 subject carries [`SessionAffinityUpdate`], which names no partition; +/// the v2 subject carries the shared [`AffinityBindingEvent`]. Both are +/// published and applied for one release so mixed-version frontends stay +/// converged. +// Compatibility with v1.6 frontends during v1.7 rolling upgrades. +// TODO(v1.8): Publish and subscribe on `SESSION_AFFINITY_SUBJECT_V2` only. pub(super) const SESSION_AFFINITY_SUBJECT: &str = "session_affinity_events"; +pub(super) const SESSION_AFFINITY_SUBJECT_V2: &str = "session_affinity_events_v2"; const OUTBOUND_CHANNEL_CAPACITY: usize = 4_096; const DIRECT_ZMQ_RCVHWM: i32 = 1_024; +/// The v1 wire payload: one binding without its partition. A receiver +/// assumes its own partition, which is what every v1.6 frontend did. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub(super) struct SessionAffinityUpdate { pub session_id: String, @@ -39,70 +52,48 @@ pub(super) struct SessionAffinityUpdate { pub writer_id: u64, } -#[derive(Clone)] -struct ReplicaUpdateSender { - tx: mpsc::Sender, -} +impl SessionAffinityUpdate { + fn from_event(event: &AffinityBindingEvent) -> Self { + Self { + session_id: event.session_id.clone(), + worker_id: event.worker_id, + dp_rank: event.dp_rank, + sequence: event.sequence, + writer_id: event.writer_id, + } + } -impl AffinityReplicaSink for ReplicaUpdateSender { - fn publish(&self, session_id: &str, target: AffinityTarget, version: AffinityVersion) { - let update = SessionAffinityUpdate { - session_id: session_id.to_string(), - worker_id: target.worker_id, - dp_rank: target.dp_rank, - sequence: version.sequence, - writer_id: version.writer_id, - }; - if let Err(error) = self.tx.try_send(update) { - tracing::trace!( - worker_id = target.worker_id, - dp_rank = ?target.dp_rank, - %error, - "dropping best-effort session affinity update" - ); + fn into_event(self, partition: &RoutingPartitionId) -> AffinityBindingEvent { + AffinityBindingEvent { + partition: partition.clone(), + session_id: self.session_id, + worker_id: self.worker_id, + dp_rank: self.dp_rank, + sequence: self.sequence, + writer_id: self.writer_id, } } } +/// Feeds received bindings to the resolver without keeping it alive. #[derive(Clone)] struct ReplicaUpdateApplier { - local_publisher_id: u64, - discovered_instances: watch::Receiver>, - table: WeakSessionAffinity, + partition: RoutingPartitionId, + resolver: Weak, } impl ReplicaUpdateApplier { - fn apply(&self, source_publisher_id: u64, update: SessionAffinityUpdate) -> bool { - if source_publisher_id == self.local_publisher_id { - return true; - } - - let Some(table) = self.table.upgrade() else { + /// `false` once the resolver is gone: the subscriber stops. + fn apply(&self, event: AffinityBindingEvent) -> bool { + let Some(resolver) = self.resolver.upgrade() else { return false; }; - table.observe_replica_sequence(update.sequence); - if !self - .discovered_instances - .borrow() - .iter() - .any(|instance| instance.id() == update.worker_id) - { - return true; - } - let target = AffinityTarget::new(update.worker_id, update.dp_rank); - let outcome = table.apply_replica_update( - update.session_id, - target, - AffinityVersion { - sequence: update.sequence, - writer_id: update.writer_id, - }, - ); - drop(table); + let (worker_id, dp_rank) = (event.worker_id, event.dp_rank); + let disposition = resolver.apply_replica_event(&self.partition, event); tracing::trace!( - worker_id = target.worker_id, - dp_rank = ?target.dp_rank, - ?outcome, + worker_id, + ?dp_rank, + ?disposition, "processed best-effort session affinity update" ); true @@ -110,126 +101,90 @@ impl ReplicaUpdateApplier { } pub(super) struct ReplicaSyncRuntime { - sender: ReplicaUpdateSender, + sink: Arc, cancel: CancellationToken, publisher_task: Option>, - subscriber_task: Option>, + subscriber_tasks: Vec>, } impl ReplicaSyncRuntime { /// Returns the runtime and this replica's writer id (its discovery instance id). - pub(super) async fn start(client: Client, table: WeakSessionAffinity) -> Result<(Self, u64)> { + pub(super) async fn start( + client: Client, + resolver: Weak, + partition: RoutingPartitionId, + ) -> Result<(Self, u64)> { let endpoint = &client.endpoint; let router_id = endpoint.drt().discovery().instance_id(); let transport_kind = endpoint.drt().default_event_transport_kind(); - let publisher = EventPublisher::for_endpoint_with_transport( + let publisher_v1 = EventPublisher::for_endpoint_with_transport( endpoint, SESSION_AFFINITY_SUBJECT, transport_kind, ) .await .context("create session affinity event publisher")?; - let publisher_id = publisher.publisher_id(); + let publisher_v2 = EventPublisher::for_endpoint_with_transport( + endpoint, + SESSION_AFFINITY_SUBJECT_V2, + transport_kind, + ) + .await + .context("create session affinity v2 event publisher")?; let applier = ReplicaUpdateApplier { - local_publisher_id: publisher_id, - discovered_instances: client.instance_source.as_ref().clone(), - table, + partition: partition.clone(), + resolver, }; let cancel = CancellationToken::new(); - let subscriber_task = - if should_use_direct_sync(transport_kind, uses_direct_zmq(transport_kind)) { - let codec = Codec::default(); - let handler_applier = applier.clone(); - let handler = move |envelope: ValidatedEnvelope| { - let source_publisher_id = envelope.publisher_id; - let update = codec - .decode_payload::(&envelope.payload) - .context("decode session affinity update")?; - handler_applier.apply(source_publisher_id, update); - Ok(()) - }; - let observer = |observation: FanInObservation| match observation.event { - FanInEvent::SequenceGap { missing } => tracing::warn!( - publisher_id = observation.publisher_id, - generation = observation.generation, - missing, - "session affinity direct-ZMQ source skipped envelopes" - ), - FanInEvent::OutOfOrder => tracing::warn!( - publisher_id = observation.publisher_id, - generation = observation.generation, - "session affinity direct-ZMQ source received a non-increasing sequence" - ), - _ => {} - }; - start_direct_zmq_fan_in( - endpoint.clone(), - SESSION_AFFINITY_SUBJECT, - DIRECT_ZMQ_RCVHWM, - Some(publisher_id), - ContinuityMode::TrackFromZero, - cancel.clone(), - handler, - observer, - ) - .await - .context("start direct-ZMQ session affinity subscriber")? - } else { - let mut subscriber = EventSubscriber::for_endpoint_with_transport( - endpoint, - SESSION_AFFINITY_SUBJECT, - transport_kind, - ) - .await - .context("create session affinity event subscriber")? - .typed::(); - let subscriber_cancel = cancel.clone(); - tokio::spawn(async move { - loop { - let event = tokio::select! { - _ = subscriber_cancel.cancelled() => return, - event = subscriber.next() => event, - }; - let Some(event) = event else { - return; - }; - let (source_publisher_id, update) = match event { - Ok((envelope, update)) => (envelope.publisher_id, update), - Err(error) => { - tracing::trace!( - %error, - "failed to receive best-effort session affinity update" - ); - continue; - } - }; - if !applier.apply(source_publisher_id, update) { - return; - } - } - }) - }; - - let (tx, mut rx) = mpsc::channel(OUTBOUND_CHANNEL_CAPACITY); - let sender = ReplicaUpdateSender { tx }; + let v2_applier = applier.clone(); + let v2_task = start_subscriber::( + &client, + SESSION_AFFINITY_SUBJECT_V2, + transport_kind, + publisher_v2.publisher_id(), + cancel.clone(), + move |event| v2_applier.apply(event), + ) + .await?; + let v1_applier = applier; + let v1_partition = partition.clone(); + let v1_task = start_subscriber::( + &client, + SESSION_AFFINITY_SUBJECT, + transport_kind, + publisher_v1.publisher_id(), + cancel.clone(), + move |update| v1_applier.apply(update.into_event(&v1_partition)), + ) + .await?; + let (tx, mut rx) = mpsc::channel::(OUTBOUND_CHANNEL_CAPACITY); let publisher_cancel = cancel.clone(); let publisher_task = tokio::spawn(async move { loop { - let update = tokio::select! { + let event = tokio::select! { _ = publisher_cancel.cancelled() => return, - update = rx.recv() => update, + event = rx.recv() => event, }; - let Some(update) = update else { + let Some(event) = event else { return; }; - if let Err(error) = publisher.publish(&update).await { + let update = SessionAffinityUpdate::from_event(&event); + if let Err(error) = publisher_v2.publish(&event).await { + tracing::trace!( + worker_id = event.worker_id, + dp_rank = ?event.dp_rank, + %error, + "failed to publish best-effort session affinity update" + ); + } + if let Err(error) = publisher_v1.publish(&update).await { tracing::trace!( worker_id = update.worker_id, dp_rank = ?update.dp_rank, %error, - "failed to publish best-effort session affinity update" + "failed to publish best-effort session affinity v1 update" ); } } @@ -237,10 +192,10 @@ impl ReplicaSyncRuntime { Ok(( Self { - sender, + sink: event_sink(partition, tx), cancel, publisher_task: Some(publisher_task), - subscriber_task: Some(subscriber_task), + subscriber_tasks: vec![v2_task, v1_task], }, router_id, )) @@ -248,17 +203,7 @@ impl ReplicaSyncRuntime { /// The sink the table publishes into. pub(super) fn sink(&self) -> Arc { - Arc::new(self.sender.clone()) - } - - #[cfg(test)] - pub(super) fn publish( - &self, - session_id: &str, - target: AffinityTarget, - version: AffinityVersion, - ) { - self.sender.publish(session_id, target, version); + Arc::clone(&self.sink) } pub(super) fn shutdown_now(&mut self) { @@ -266,24 +211,120 @@ impl ReplicaSyncRuntime { if let Some(task) = self.publisher_task.take() { task.abort(); } - drop(self.subscriber_task.take()); + self.subscriber_tasks.clear(); } #[cfg(test)] - pub(super) fn for_test(capacity: usize) -> (Self, mpsc::Receiver) { + pub(super) fn for_test(capacity: usize) -> (Self, mpsc::Receiver) { let (tx, rx) = mpsc::channel(capacity); ( Self { - sender: ReplicaUpdateSender { tx }, + sink: event_sink(crate::kv_router::embedded::embedded_partition_key(None), tx), cancel: CancellationToken::new(), publisher_task: None, - subscriber_task: None, + subscriber_tasks: Vec::new(), }, rx, ) } } +/// A sink that queues each binding for the publisher task; a full queue +/// drops the update (replication is best effort). +fn event_sink( + partition: RoutingPartitionId, + tx: mpsc::Sender, +) -> Arc { + Arc::new(AffinityEventSink::new(partition, move |event| { + if let Err(error) = tx.try_send(event) { + tracing::trace!(%error, "dropping best-effort session affinity update"); + } + })) +} + +/// Subscribe to `subject` on the deployment's event transport and feed each +/// decoded payload to `apply`, which returns `false` to stop. +async fn start_subscriber( + client: &Client, + subject: &'static str, + transport_kind: EventTransportKind, + publisher_id: u64, + cancel: CancellationToken, + apply: impl Fn(T) -> bool + Clone + Send + Sync + 'static, +) -> Result> +where + T: DeserializeOwned + Send + 'static, +{ + let endpoint = &client.endpoint; + if should_use_direct_sync(transport_kind, uses_direct_zmq(transport_kind)) { + let codec = Codec::default(); + let handler = move |envelope: ValidatedEnvelope| { + let payload = codec + .decode_payload::(&envelope.payload) + .context("decode session affinity update")?; + apply(payload); + Ok(()) + }; + let observer = move |observation: FanInObservation| match observation.event { + FanInEvent::SequenceGap { missing } => tracing::warn!( + subject, + publisher_id = observation.publisher_id, + generation = observation.generation, + missing, + "session affinity direct-ZMQ source skipped envelopes" + ), + FanInEvent::OutOfOrder => tracing::warn!( + subject, + publisher_id = observation.publisher_id, + generation = observation.generation, + "session affinity direct-ZMQ source received a non-increasing sequence" + ), + _ => {} + }; + return start_direct_zmq_fan_in( + endpoint.clone(), + subject, + DIRECT_ZMQ_RCVHWM, + Some(publisher_id), + ContinuityMode::TrackFromZero, + cancel, + handler, + observer, + ) + .await + .context("start direct-ZMQ session affinity subscriber"); + } + let mut subscriber = + EventSubscriber::for_endpoint_with_transport(endpoint, subject, transport_kind) + .await + .context("create session affinity event subscriber")? + .typed::(); + Ok(tokio::spawn(async move { + loop { + let event = tokio::select! { + _ = cancel.cancelled() => return, + event = subscriber.next() => event, + }; + let Some(event) = event else { + return; + }; + let payload = match event { + Ok((_envelope, payload)) => payload, + Err(error) => { + tracing::trace!( + %error, + "failed to receive best-effort session affinity update" + ); + continue; + } + }; + if !apply(payload) { + return; + } + } + })) +} + impl Drop for ReplicaSyncRuntime { fn drop(&mut self) { self.shutdown_now(); @@ -297,7 +338,11 @@ fn should_use_direct_sync(transport_kind: EventTransportKind, direct_zmq_topolog #[cfg(test)] mod tests { use super::*; - use crate::session_affinity::{AlwaysLive, HostAffinity}; + use crate::session_affinity::{AlwaysLive, HostAffinity, LiveWorkers}; + use dynamo_kv_router::services::selection::affinity::{ + AffinityTarget, AffinityVersion, ReplicaApplyOutcome, ReplicaEventDisposition, + SessionAffinity, SessionAffinityConfig, + }; use dynamo_runtime::{ DistributedRuntime, Runtime, discovery::{DiscoveryQuery, EventChannelQuery}, @@ -315,7 +360,8 @@ mod tests { #[tokio::test] async fn replica_update_backpressure_is_nonfatal() { let (runtime, mut rx) = ReplicaSyncRuntime::for_test(1); - runtime.publish( + let sink = runtime.sink(); + sink.publish( "first", AffinityTarget { worker_id: 10, @@ -326,7 +372,7 @@ mod tests { writer_id: 7, }, ); - runtime.publish( + sink.publish( "second", AffinityTarget { worker_id: 11, @@ -341,9 +387,102 @@ mod tests { let update = rx.recv().await.unwrap(); assert_eq!(update.session_id, "first"); assert_eq!(update.writer_id, 7); + assert_eq!( + update.partition, + crate::kv_router::embedded::embedded_partition_key(None) + ); assert!(rx.try_recv().is_err()); } + /// A v1 update (no partition) and a v2 event both apply to the receiving + /// partition; a v2 event for another partition does not. + #[tokio::test] + async fn v1_and_v2_updates_both_apply_to_the_partition() { + let partition = crate::kv_router::embedded::embedded_partition_key(Some("model")); + let table = + SessionAffinity::with_config(SessionAffinityConfig::new(Duration::from_secs(10))) + .unwrap(); + let resolver = Arc::new(AffinityResolver::new( + table, + Arc::new(LiveWorkers([10, 11].into_iter().collect())), + )); + let applier = ReplicaUpdateApplier { + partition: partition.clone(), + resolver: Arc::downgrade(&resolver), + }; + + let v1 = SessionAffinityUpdate { + session_id: "legacy".to_string(), + worker_id: 10, + dp_rank: Some(0), + sequence: 5, + writer_id: 9, + }; + assert!(applier.apply(v1.into_event(&partition))); + assert_eq!( + resolver.table().query_target("legacy", None).unwrap(), + Some(AffinityTarget::new(10, Some(0))) + ); + + let v2 = AffinityBindingEvent { + partition: partition.clone(), + session_id: "current".to_string(), + worker_id: 11, + dp_rank: None, + sequence: 6, + writer_id: 9, + }; + assert!(applier.apply(v2)); + assert_eq!( + resolver.table().query_target("current", None).unwrap(), + Some(AffinityTarget::new(11, None)) + ); + + let elsewhere = AffinityBindingEvent { + partition: crate::kv_router::embedded::embedded_partition_key(Some("other")), + session_id: "elsewhere".to_string(), + worker_id: 10, + dp_rank: None, + sequence: 7, + writer_id: 9, + }; + assert_eq!( + resolver.apply_replica_event(&partition, elsewhere), + ReplicaEventDisposition::OtherPartition + ); + assert_eq!( + resolver.table().query_target("elsewhere", None).unwrap(), + None + ); + + // The same binding on both subjects is idempotent. + let again = SessionAffinityUpdate { + session_id: "current".to_string(), + worker_id: 11, + dp_rank: None, + sequence: 6, + writer_id: 9, + }; + assert_eq!( + resolver.apply_replica_event(&partition, again.into_event(&partition)), + ReplicaEventDisposition::Applied(ReplicaApplyOutcome::Refreshed) + ); + + drop(resolver); + assert!( + !applier.apply( + SessionAffinityUpdate { + session_id: "late".to_string(), + worker_id: 10, + dp_rank: None, + sequence: 8, + writer_id: 9, + } + .into_event(&partition) + ) + ); + } + fn host_affinity() -> HostAffinity { HostAffinity::standalone( Duration::from_secs(10), @@ -355,27 +494,34 @@ mod tests { #[tokio::test] async fn rejected_worker_update_still_advances_replica_clock() { - let affinity = host_affinity(); - let table = affinity.resolver().table(); + let partition = crate::kv_router::embedded::embedded_partition_key(None); + let table = + SessionAffinity::with_config(SessionAffinityConfig::new(Duration::from_secs(10))) + .unwrap(); + let resolver = Arc::new(AffinityResolver::new( + table, + Arc::new(LiveWorkers(Default::default())), + )); + let table = resolver.table(); let baseline = table.next_version(); let sequence = baseline.sequence.saturating_add(10); - let (_tx, discovered_instances) = watch::channel(Vec::new()); let applier = ReplicaUpdateApplier { - local_publisher_id: 7, - discovered_instances, - table: table.downgrade(), + partition: partition.clone(), + resolver: Arc::downgrade(&resolver), }; - assert!(applier.apply( - 8, - SessionAffinityUpdate { - session_id: "unknown-worker".to_string(), - worker_id: 10, - dp_rank: Some(0), - sequence, - writer_id: 9, - } - )); + assert!( + applier.apply( + SessionAffinityUpdate { + session_id: "unknown-worker".to_string(), + worker_id: 10, + dp_rank: Some(0), + sequence, + writer_id: 9, + } + .into_event(&partition) + ) + ); assert_eq!(table.next_version().sequence, sequence.saturating_add(1)); } diff --git a/lib/llm/src/session_affinity/tests.rs b/lib/llm/src/session_affinity/tests.rs index c7ea3310eddd..17698d22be82 100644 --- a/lib/llm/src/session_affinity/tests.rs +++ b/lib/llm/src/session_affinity/tests.rs @@ -22,9 +22,8 @@ use tokio_util::sync::CancellationToken; use super::SessionAffinityMode::{Hard, Soft}; use super::{ AffinityTarget, AlwaysLive, Hold, HostAffinity, LlmResponse, affinity_error, affinity_id, - explicit_target, from_table, - replica_sync::{ReplicaSyncRuntime, SessionAffinityUpdate}, - subagent_group_affinity_id, to_table, tracked_stream, + explicit_target, from_table, replica_sync::ReplicaSyncRuntime, subagent_group_affinity_id, + to_table, tracked_stream, }; use crate::{ preprocessor::PreprocessedRequest, @@ -37,8 +36,8 @@ use crate::{ types::Annotated, }; use dynamo_kv_router::services::selection::affinity::{ - AffinityResolver, AffinityVersion, MAX_SESSION_AFFINITY_ID_BYTES, ReplicaApplyOutcome, - Resolution, SessionAffinity, SessionAffinityConfig, + AffinityBindingEvent, AffinityResolver, AffinityVersion, MAX_SESSION_AFFINITY_ID_BYTES, + ReplicaApplyOutcome, Resolution, SessionAffinity, SessionAffinityConfig, }; fn session_id() -> SessionAffinityId { @@ -181,7 +180,7 @@ impl TestAffinity { &self, router_id: u64, capacity: usize, - ) -> mpsc::Receiver { + ) -> mpsc::Receiver { let (replica, rx) = ReplicaSyncRuntime::for_test(capacity); assert!( self.table().enable_replication(router_id, replica.sink()),