From 4d118d834df232ad415a01d061a45895f9f3f683 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Mon, 29 Jun 2026 21:54:01 -0700 Subject: [PATCH 1/6] feat(router): coordinate session affinity across replicas Signed-off-by: PeaBrane --- .../configuration/groups/router_args.py | 8 +- docs/agents/session-ids.md | 19 +- docs/components/frontend/nvext.md | 12 +- .../components/router/router-configuration.md | 62 +- lib/llm/src/kv_router/push_router.rs | 207 +++---- .../kv_router/push_router/request_guard.rs | 3 +- .../src/kv_router/push_router/selection.rs | 33 +- lib/llm/src/session_affinity/coordinator.rs | 472 ++++++++++----- lib/llm/src/session_affinity/mod.rs | 7 +- lib/llm/src/session_affinity/push_router.rs | 256 ++++----- lib/llm/src/session_affinity/tests.rs | 535 +++++++++++++----- lib/runtime/src/discovery/kube.rs | 32 +- lib/runtime/src/discovery/kv_store.rs | 510 ++++++++++++++++- lib/runtime/src/discovery/mod.rs | 52 ++ .../pipeline/network/egress/push_router.rs | 43 ++ lib/runtime/src/storage/kv.rs | 4 + tests/router/common.py | 203 +++++++ tests/router/router_process.py | 8 +- tests/router/test_router_e2e_with_mockers.py | 39 ++ 19 files changed, 1896 insertions(+), 609 deletions(-) diff --git a/components/src/dynamo/common/configuration/groups/router_args.py b/components/src/dynamo/common/configuration/groups/router_args.py index f4bb611b8769..1c9b334bcd94 100644 --- a/components/src/dynamo/common/configuration/groups/router_args.py +++ b/components/src/dynamo/common/configuration/groups/router_args.py @@ -244,9 +244,11 @@ def add_arguments(self, parser) -> None: env_var="DYN_ROUTER_SESSION_AFFINITY_TTL_SECS", default=None, help=( - "Enable router-local session affinity with this idle TTL in seconds. " - "Affinity is disabled when this option is omitted. " - "This is independent of KV prediction TTL settings." + "Enable session affinity and set the process-local cache eviction TTL " + "in seconds. etcd and shared FileStore use immutable distributed claims " + "whose lifetime follows the creating frontend, not this TTL. Memory and " + "Kubernetes discovery remain process-local. Affinity is disabled when " + "this option is omitted." ), arg_type=int, dest="session_affinity_ttl_secs", diff --git a/docs/agents/session-ids.md b/docs/agents/session-ids.md index 6efb096f1ab3..c8aa3b4d36c3 100644 --- a/docs/agents/session-ids.md +++ b/docs/agents/session-ids.md @@ -7,7 +7,11 @@ subtitle: Identify agent sessions from supported coding agents and custom client A session ID is the stable identifier Dynamo uses for one agent reasoning/tool chain. A root agent, planner, researcher subagent, or OpenCode subtask can each have its own session. Every LLM request in that chain should carry the same `session_id`; child sessions can also carry a `parent_session_id` so traces and replay tools can rebuild the tree. Some academic papers also call this a `program_id`. -Session identity is passive metadata. Sending `X-Dynamo-Session-ID` does not enable sticky sessions or change request placement. Tracing records the identity when `DYN_REQUEST_TRACE` is enabled, and a session-aware routing policy can consume it only when that policy is configured separately. +Session identity is passive metadata unless session affinity is explicitly enabled. +Sending `X-Dynamo-Session-ID` alone does not change request placement. Tracing records +the identity when `DYN_REQUEST_TRACE` is enabled. When +`--router-session-affinity-ttl-secs` is configured, the router uses the ID for an +immutable endpoint- and phase-scoped worker binding. ## Session ID Inputs @@ -29,7 +33,18 @@ Dynamo also recognizes the current stable identity headers emitted by the follow | Codex | `session-id` | None | `session-id` becomes the `session_id`. | | OpenCode | `x-session-id` | `x-parent-session-id` | `x-session-id` becomes the `session_id`; `x-parent-session-id` becomes `parent_session_id` when present. | -`X-Dynamo-Session-Final` applies with either canonical or agent-native session identity. +`X-Dynamo-Session-Final` applies with either canonical or agent-native session +identity. With session affinity enabled, a final request routes normally and then +terminally closes its binding. Close invalidation across replicas is eventual. Do not +send more requests with that session ID after close. + +etcd and FileStore on a shared filesystem coordinate bindings across frontend +processes. MemoryStore and Kubernetes discovery retain process-local affinity only. +The affinity TTL controls local cache cleanup, not the distributed claim lifetime. +Claims follow the creating frontend's etcd lease or FileStore ownership. If a claim +expires or its bound worker disappears, create a new session with a new ID instead of +reusing or rebinding the old ID. See [Router session affinity](../components/router/router-configuration.md#session-affinity) +for the full contract. ### Custom Agent Harnesses diff --git a/docs/components/frontend/nvext.md b/docs/components/frontend/nvext.md index 0970e1ad9524..ad7b41457537 100644 --- a/docs/components/frontend/nvext.md +++ b/docs/components/frontend/nvext.md @@ -54,7 +54,7 @@ token IDs, pass integer IDs in the normal `stop` array, for example `"stop": [576]`. Strings such as `"token_id:576"` remain literal string stop sequences and are not parsed as token IDs. -### Header Overrides +### Header overrides Routing fields can also be set via HTTP headers, which take priority over `nvext` values: @@ -75,9 +75,13 @@ session headers described in [Session IDs](../../agents/session-ids.md); `nvext` does not accept session identity fields. When session affinity is enabled with `--router-session-affinity-ttl-secs`, the -router also uses `X-Dynamo-Session-ID` for router-local affinity. See -[Configuration and Tuning](../router/router-configuration.md#session-affinity) -for routing behavior and TTL settings. +router uses `X-Dynamo-Session-ID` for immutable endpoint- and phase-scoped affinity. +On etcd and shared FileStore, replicas coordinate through a distributed claim while +the request hot path uses a process-local cache. Existing local or shared bindings +override routing headers; the headers above are proposals only when no binding exists. +Memory and Kubernetes discovery do not provide cross-process affinity. See +[Configuration and Tuning](../router/router-configuration.md#session-affinity) for +claim lifetime, cache TTL, terminal close, and failure behavior. For trace sink configuration and JSONL schema details, see [Agent Tracing](../../agents/agent-tracing.md). diff --git a/docs/components/router/router-configuration.md b/docs/components/router/router-configuration.md index 09e8e13a9aa9..17be76ca6bb6 100644 --- a/docs/components/router/router-configuration.md +++ b/docs/components/router/router-configuration.md @@ -111,31 +111,43 @@ a value from `1` through `31536000` to enable it, then send `X-Dynamo-Session-ID` to keep related requests on one worker. Supplying the header without the TTL option provides session identity but does not enable router affinity. -The first successfully dispatched request binds the session ID to its selected -worker and, when available, data-parallel rank. Later requests exact-dispatch to -that target without transport fallback. Concurrent requests can share a binding. -Active requests prevent expiry. When a request lease ends after EOF, early drop, -error, or cancellation, the idle timer restarts. A missing bound worker or a -non-cancellation selection, setup, dispatch, or target-validation failure invalidates -the binding. - -The configured value is the idle timeout. It is independent of -`--router-ttl-secs` and `--router-predicted-ttl-secs`. Omit the session-affinity -option to keep affinity disabled. - -If the bound worker disappears, Dynamo invalidates the binding so a subsequent -selection can bind an available worker. Router restart clears all bindings. Bindings -are not shared between frontend replicas. - -Direct mode still requires the phase-appropriate explicit worker ID on every -affinity request. The stored binding validates that target but does not supply a -missing ID. In disaggregated serving, prefill and decode use separate phase-local -bindings. If no prefill router is active, only the decode or aggregated binding is -created. - -Session affinity does not create a backend session or send lifecycle RPCs. There is -no explicit unbind; idle expiry removes only router-local state. The same session -ID is available to tracing and other explicitly configured consumers. +The first affinity request creates one immutable binding from the session ID to a +worker and, when available, a data-parallel rank. The binding is scoped to the +existing endpoint and phase, so disaggregated prefill and decode routes remain +separate. Later requests exact-dispatch to that target without transport fallback. +An existing local or shared binding takes precedence over explicit routing headers; +those headers are proposals only while the claim is absent. Direct mode therefore +requires an explicit target for a new binding, but an existing binding supplies the +target for later requests. Query-only requests remain read-only and do not create or +close claims. + +With etcd or FileStore on a filesystem shared by all replicas, frontends coordinate +through an immutable distributed claim. The existing-session hot path reads only the +process-local cache. A cache miss reads shared storage first and attempts an atomic +insertion only when the claim is absent. Racing frontends all cache and dispatch to +the stored winner. Storage errors fail the request before scheduler bookkeeping or +dispatch. MemoryStore coordinates only callers sharing the same process and store. +Kubernetes discovery does not provide cross-process affinity and keeps process-local +behavior. + +For distributed backends, `--router-session-affinity-ttl-secs` controls only +process-local cache eviction. A cache miss after local eviction reloads the immutable +claim. The claim itself follows the creating frontend's existing etcd lease or +FileStore ownership lifetime; it is not a global idle-session timeout. Delete events +eventually invalidate other frontend caches. Watch lag, disconnect, or restart clears +the entire local affinity cache, and later requests reload claims on demand. + +`X-Dynamo-Session-Final: true` marks a terminal request. Dynamo routes that request +normally, then evicts the closing frontend's cache entry and idempotently deletes the +shared claim. Other replicas observe the delete eventually. Close must not race active +requests, and callers must not use that session ID again. The same no-reuse rule +applies after claim expiry. If the bound worker disappears while the claim exists, +exact dispatch fails; start a new session with a new session ID. + +Global idle-session TTL, rebinding, dead-worker replacement, compare-and-swap updates, +fencing, generations, broader `WorkerSet` affinity, and backend-tokenized path +expansion are outside this contract. The setting remains independent of +`--router-ttl-secs` and `--router-predicted-ttl-secs`; omit it to disable affinity. ### AIC Prefill Load Model diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index 0c194ece4f10..33f43326127a 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -5,13 +5,14 @@ use std::{sync::Arc, time::Duration}; use dynamo_kv_router::protocols::{TokensWithHashes, WorkerWithDpRank}; use dynamo_runtime::{ - error::{ErrorType, match_error_chain}, + discovery::ClaimPayloadFuture, metrics::frontend_perf::{STAGE_ROUTE, StageGuard}, pipeline::{ AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, PushRouter, ResponseStream, SingleIn, async_trait, }, protocols::annotated::Annotated, + traits::DistributedRuntimeProvider, }; use futures::stream::{self, StreamExt}; use tracing::Instrument; @@ -24,7 +25,7 @@ use crate::{ timing::{RequestPhase, RoutingData}, }, session_affinity::{ - AffinityAcquire, AffinityCoordinator, AffinityTarget, affinity_id, explicit_target, + AffinityCoordinator, AffinityTarget, ResolvedAffinity, affinity_id, session_final, }, }; @@ -52,7 +53,13 @@ impl KvPushRouter { session_affinity_ttl: Option, ) -> Result { let affinity = session_affinity_ttl - .map(AffinityCoordinator::new) + .map(|ttl| { + AffinityCoordinator::new_distributed( + ttl, + inner.client.endpoint.id().to_string(), + inner.client.endpoint.drt().discovery(), + ) + }) .transpose()?; // Eagerly register router request metrics (as zeros) so they are @@ -121,7 +128,7 @@ impl KvPushRouter { request: &SingleIn, phase: RequestPhase, is_query_only: bool, - ) -> Result<(WorkerSelection, Option), Error> { + ) -> Result<(WorkerSelection, Option), Error> { let Some(affinity) = self.affinity.as_ref() else { return Ok(( self.select_request(request, phase, is_query_only, None) @@ -136,9 +143,8 @@ impl KvPushRouter { None, )); }; - let explicit = explicit_target(request, phase)?; if is_query_only { - let target = affinity.query_target(&session_id, explicit)?; + let target = affinity.query_target(&session_id)?; let worker = target.and_then(affinity_worker); return Ok(( self.select_request(request, phase, true, worker).await?, @@ -148,32 +154,20 @@ impl KvPushRouter { let request_context = request.context(); let operation = affinity - .acquire_with_context(&session_id, explicit, request_context.as_ref()) + .acquire_with_context(&session_id, request_context.as_ref()) .await?; - let worker = operation.target().and_then(affinity_worker); - match self.select_request(request, phase, false, worker).await { - Ok(selection) => Ok((selection, Some(operation))), - Err(error) if match_error_chain(error.as_ref(), &[ErrorType::Cancelled], &[]) => { - Err(error) - } - Err(_) if operation.target().is_some() && explicit.is_none() => { - operation.invalidate(); - let retry = affinity - .acquire_with_context(&session_id, None, request_context.as_ref()) - .await?; - match self.select_request(request, phase, false, None).await { - Ok(selection) => Ok((selection, Some(retry))), - Err(retry_error) => { - retry.invalidate(); - Err(retry_error) - } - } - } - Err(error) => { - operation.invalidate(); - Err(error) - } - } + let proposed_payload: ClaimPayloadFuture<'_> = Box::pin(async { + let selection = self.select_request(request, phase, true, None).await?; + let target = AffinityTarget { + worker_id: selection.instance_id, + dp_rank: Some(selection.dp_rank), + }; + Ok(serde_json::to_value(target)?) + }); + let resolved = operation.resolve(proposed_payload).await?; + let worker = affinity_worker(resolved.target()); + let selection = self.select_request(request, phase, false, worker).await?; + Ok((selection, Some(resolved))) } async fn track_selection( @@ -394,52 +388,28 @@ impl KvPushRouter { let phase_label = phase.to_string(); let route_guard = StageGuard::new(STAGE_ROUTE, &phase_label); let is_query_only = request.get_annotation_value("query_instance_id").is_some(); - let (mut selection, mut operation) = self + let close_on_finish = !is_query_only && session_final(request.content()); + let (mut selection, operation) = self .select_with_affinity(&request, phase, is_query_only) .await?; - let mut guard = match self + let mut guard = self .track_selection(&request, &mut selection, is_query_only) - .await - { - Ok(guard) => guard, - Err(error) => { - if let Some(operation) = operation.take() { - operation.invalidate(); - } - return Err(error); - } - }; + .await?; let metadata = match prepare(&mut request, selection.instance_id, Some(selection.dp_rank)) { Ok(metadata) => metadata, Err(error) => { guard.abort().await; - if let Some(operation) = operation.take() { - operation.invalidate(); - } return Err(error); } }; - let selected_target = AffinityTarget { - worker_id: selection.instance_id, - dp_rank: Some(selection.dp_rank), - }; drop(route_guard); - let stream = match self + let stream = self .dispatch_selection(request, selection, guard, true) - .await - { - Ok(stream) => stream, - Err(error) => { - if let Some(operation) = operation.take() { - operation.invalidate(); - } - return Err(error); - } - }; + .await?; let Some(operation) = operation else { return Ok((metadata, stream)); }; - Ok((metadata, operation.into_stream(selected_target, stream)?)) + Ok((metadata, operation.into_stream(stream, close_on_finish))) } } @@ -478,7 +448,8 @@ impl AsyncEngine, ManyOut, ManyOut guard, - Err(error) => { - if let Some(operation) = operation.take() { - operation.invalidate(); - } - return Err(error); - } - }; + let guard = self + .track_selection(&request, &mut selection, false) + .await?; drop(route_guard); - let selected_target = AffinityTarget { - worker_id: selection.instance_id, - dp_rank: Some(selection.dp_rank), - }; - let stream = match self + let stream = self .dispatch_selection(request, selection, guard, operation.is_some()) - .await - { - Ok(stream) => stream, - Err(error) => { - if let Some(operation) = operation.take() { - operation.invalidate(); - } - return Err(error); - } - }; + .await?; match operation { - Some(operation) => operation.into_stream(selected_target, stream), + Some(operation) => Ok(operation.into_stream(stream, close_on_finish)), None => Ok(stream), } } @@ -653,7 +605,10 @@ mod tests { .unwrap(); let endpoint = component.endpoint("generate"); let client = endpoint.client().await.unwrap(); - let workers = HashMap::from([(7, ModelRuntimeConfig::default())]); + let workers = HashMap::from([ + (7, ModelRuntimeConfig::default()), + (8, ModelRuntimeConfig::default()), + ]); let (_tx, workers) = watch::channel(workers); let config = KvRouterConfig { skip_initial_worker_wait: true, @@ -701,17 +656,19 @@ mod tests { worker_id: 7, dp_rank: Some(0), }; - let AffinityAcquire::Initialize(initializer) = router + let resolved = router .affinity .as_ref() .unwrap() - .acquire(&session_id, None) + .acquire(&session_id) .await .unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(original_target).unwrap()); + .resolve(Box::pin(async move { + Ok(serde_json::to_value(original_target)?) + })) + .await + .unwrap(); + drop(resolved); let controller = Controller::new("cancelled-selection-request".to_string()); controller.stop(); @@ -734,23 +691,69 @@ mod tests { .affinity .as_ref() .unwrap() - .query_target(&session_id, None) + .query_target(&session_id) .unwrap(), Some(original_target) ); - let AffinityAcquire::Bound { target, lease } = router + let resolved = router .affinity .as_ref() .unwrap() - .acquire(&session_id, None) + .acquire(&session_id) .await .unwrap() - else { - panic!("cancellation must preserve the existing binding"); + .resolve(Box::pin(async move { + Ok(serde_json::to_value(AffinityTarget { + worker_id: 8, + dp_rank: Some(0), + })?) + })) + .await + .unwrap(); + assert_eq!(resolved.target(), original_target); + + drop(router); + runtime.shutdown(); + } + + #[tokio::test] + async fn session_affinity_binding_overrides_conflicting_explicit_proposal() { + let (router, runtime) = router(Some(Duration::from_secs(10))).await; + let session_id = SessionAffinityId::new("conflicting-explicit-proposal"); + let bound_target = AffinityTarget { + worker_id: 7, + dp_rank: Some(0), }; - assert_eq!(target, original_target); - drop(lease); + let resolved = router + .affinity + .as_ref() + .unwrap() + .acquire(&session_id) + .await + .unwrap() + .resolve(Box::pin( + async move { Ok(serde_json::to_value(bound_target)?) }, + )) + .await + .unwrap(); + drop(resolved); + + let mut content = request(); + content.routing_mut().backend_instance_id = Some(8); + content.routing_mut().decode_worker_id = Some(8); + content.routing_mut().dp_rank = Some(0); + let mut request = Context::new(content); + request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id); + + let (selection, resolved) = router + .select_with_affinity(&request, RequestPhase::Aggregated, false) + .await + .unwrap(); + assert_eq!(selection.instance_id, 7); + assert_eq!(selection.dp_rank, 0); + assert_eq!(resolved.unwrap().target(), bound_target); + router.chooser.free(request.context().id()).await.unwrap(); drop(router); runtime.shutdown(); diff --git a/lib/llm/src/kv_router/push_router/request_guard.rs b/lib/llm/src/kv_router/push_router/request_guard.rs index 6f0df3cecf15..a944b22d0fc4 100644 --- a/lib/llm/src/kv_router/push_router/request_guard.rs +++ b/lib/llm/src/kv_router/push_router/request_guard.rs @@ -245,8 +245,7 @@ impl OutputBlockTracker { /// Coordinates scheduler cleanup, observability, and streamed load tracking. /// -/// Session-affinity lifetime is separate: `AffinityAcquire` and -/// `AffinityLease` own binding commit, release, and invalidation. +/// Session-affinity lifetime is separate: `ResolvedAffinity` owns the binding lease. pub(super) struct RequestGuard { cleanup: RequestCleanup, observability: RequestObservability, diff --git a/lib/llm/src/kv_router/push_router/selection.rs b/lib/llm/src/kv_router/push_router/selection.rs index 980c676383c0..5ac97db65182 100644 --- a/lib/llm/src/kv_router/push_router/selection.rs +++ b/lib/llm/src/kv_router/push_router/selection.rs @@ -140,9 +140,7 @@ impl KvPushRouter { let affinity_pin = options .affinity_worker .map(|worker| (worker.worker_id, Some(worker.dp_rank))); - let Some((pinned_worker_id, requested_dp_rank)) = - merge_affinity_pin(explicit_pin, affinity_pin) - else { + let Some((pinned_worker_id, requested_dp_rank)) = affinity_pin.or(explicit_pin) else { let _nvtx_kv = dynamo_nvtx_range!("route.kv_match"); let selection = self .select_best_match(BestMatchArgs { @@ -235,21 +233,6 @@ impl KvPushRouter { } } -fn merge_affinity_pin( - explicit: Option<(u64, Option)>, - affinity: Option<(u64, Option)>, -) -> Option<(u64, Option)> { - match (explicit, affinity) { - (Some((worker_id, None)), Some((affinity_worker_id, affinity_rank))) - if worker_id == affinity_worker_id => - { - Some((worker_id, affinity_rank)) - } - (Some(explicit), _) => Some(explicit), - (None, affinity) => affinity, - } -} - fn resolve_pinned_worker_rank( worker_id: WorkerId, requested_dp_rank: Option, @@ -295,7 +278,7 @@ mod tests { scheduling::{RoutingEligibility, WorkerEligibilityError}, }; - use super::{merge_affinity_pin, pinned_worker_hint, resolve_pinned_worker_rank}; + use super::{pinned_worker_hint, resolve_pinned_worker_rank}; use crate::{ local_model::runtime_config::ModelRuntimeConfig, protocols::common::{preprocessor::RoutingHints, timing::RequestPhase}, @@ -323,18 +306,6 @@ mod tests { assert!(error.contains("requires an explicit dp_rank")); } - #[test] - fn affinity_pin_supplies_rank_for_matching_explicit_worker() { - assert_eq!( - merge_affinity_pin(Some((7, None)), Some((7, Some(0)))), - Some((7, Some(0))) - ); - assert_eq!( - merge_affinity_pin(Some((7, Some(2))), Some((7, Some(3)))), - Some((7, Some(2))) - ); - } - #[test] fn pinned_worker_hint_prefill_uses_prefill_worker_before_backend() { let routing = RoutingHints { diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index 7702309b3bd1..8dea2e88f599 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -1,6 +1,21 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +//! Session affinity uses a process-local cache in front of immutable shared claims. +//! Cache hits exact-route without distributed I/O. On a cache miss, the discovery +//! backend reads the claim first and evaluates the query-only routing proposal only +//! when no claim exists. The payload returned by claim arbitration is authoritative: +//! distributed mode caches only `Created` or `Existing` payloads, and racing losers +//! discard their proposal, cache the winner, and dispatch to it. Explicit worker and +//! rank headers are proposals only while no binding exists. +//! +//! Shared claim deletion invalidates local caches eventually through `Delete` and +//! `Reset` events. The configured affinity TTL evicts only local cache entries; it +//! does not expire or replace a shared claim. Bindings are never rebound in v1. If a +//! bound worker disappears, exact dispatch fails without fallback and the caller must +//! use a new session ID. Explicit close requires no concurrent active requests, is +//! terminal, and the closed session ID must not be reused. + use std::{ pin::Pin, sync::{ @@ -13,12 +28,17 @@ use std::{ use dashmap::{DashMap, mapref::entry::Entry}; use dynamo_runtime::{ + discovery::{ClaimEvent, ClaimOutcome, ClaimPayloadFuture, Discovery}, engine::{AsyncEngineContext, AsyncEngineContextProvider}, error::{DynamoError, ErrorType}, pipeline::{Error, ManyOut, ResponseStream}, }; use futures::Stream; -use tokio::{sync::Notify, time::Instant}; +use serde::{Deserialize, Serialize}; +use tokio::{ + sync::{Notify, broadcast}, + time::Instant, +}; use tokio_util::sync::CancellationToken; use super::{ @@ -33,7 +53,7 @@ use crate::{ }, }; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, Serialize)] pub struct AffinityTarget { pub worker_id: u64, pub dp_rank: Option, @@ -54,6 +74,7 @@ enum AffinityEntry { struct AffinityCoordinatorInner { entries: DashMap, + claims: ClaimCoordination, ttl: Duration, max_entries: usize, max_session_id_bytes: usize, @@ -61,11 +82,71 @@ struct AffinityCoordinatorInner { next_revision: AtomicU64, cancel: CancellationToken, #[cfg(test)] + probe: AffinityCoordinatorProbe, +} + +#[derive(Clone)] +struct ClaimCoordination { + scope: String, + discovery: Option>, +} + +impl ClaimCoordination { + fn key(&self, session_id: &SessionAffinityId) -> String { + format!( + "{}/{}", + self.scope, + blake3::hash(session_id.as_str().as_bytes()).to_hex() + ) + } + + fn subscribe(&self) -> Option> { + self.discovery + .as_ref() + .and_then(|discovery| discovery.subscribe_claim_events()) + } + + async fn resolve( + &self, + key: &str, + proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> Result<(serde_json::Value, bool), Error> { + let Some(discovery) = self.discovery.as_ref() else { + return Ok((proposed_payload.as_mut().await?, true)); + }; + + match discovery.create_or_get_claim(key, proposed_payload).await? { + ClaimOutcome::Created(payload) => Ok((payload, true)), + ClaimOutcome::Existing(payload) => Ok((payload, false)), + ClaimOutcome::Unsupported => Ok((proposed_payload.as_mut().await?, true)), + } + } + + async fn close(&self, key: &str) -> anyhow::Result<()> { + let Some(discovery) = self.discovery.as_ref() else { + return Ok(()); + }; + discovery.close_claim(key).await?; + Ok(()) + } +} + +#[cfg(test)] +struct AffinityCoordinatorProbe { reaper_started: Arc, - #[cfg(test)] waiter_observed: Arc, } +#[cfg(test)] +impl AffinityCoordinatorProbe { + fn new() -> Self { + Self { + reaper_started: Arc::new(Notify::new()), + waiter_observed: Arc::new(Notify::new()), + } + } +} + impl Drop for AffinityCoordinatorInner { fn drop(&mut self) { self.cancel.cancel(); @@ -83,6 +164,22 @@ impl AffinityCoordinator { ttl, MAX_SESSION_AFFINITY_ENTRIES, MAX_SESSION_AFFINITY_ID_BYTES, + "local".to_string(), + None, + ) + } + + pub(crate) fn new_distributed( + ttl: Duration, + claim_scope: String, + discovery: Arc, + ) -> Result { + Self::new_with_limits( + ttl, + MAX_SESSION_AFFINITY_ENTRIES, + MAX_SESSION_AFFINITY_ID_BYTES, + claim_scope, + Some(discovery), ) } @@ -90,6 +187,8 @@ impl AffinityCoordinator { ttl: Duration, max_entries: usize, max_session_id_bytes: usize, + claim_scope: String, + discovery: Option>, ) -> Result { if !(Duration::from_secs(1)..=Duration::from_secs(MAX_SESSION_AFFINITY_TTL_SECS)) .contains(&ttl) @@ -100,6 +199,10 @@ impl AffinityCoordinator { } let inner = Arc::new(AffinityCoordinatorInner { entries: DashMap::new(), + claims: ClaimCoordination { + scope: claim_scope, + discovery, + }, ttl, max_entries, max_session_id_bytes, @@ -107,20 +210,79 @@ impl AffinityCoordinator { next_revision: AtomicU64::new(1), cancel: CancellationToken::new(), #[cfg(test)] - reaper_started: Arc::new(Notify::new()), - #[cfg(test)] - waiter_observed: Arc::new(Notify::new()), + probe: AffinityCoordinatorProbe::new(), }); Self::spawn_reaper(&inner); + Self::spawn_claim_listener(&inner); Ok(Self { inner }) } + fn spawn_claim_listener(inner: &Arc) { + let Some(mut events) = inner.claims.subscribe() else { + return; + }; + let weak = Arc::downgrade(inner); + let cancel = inner.cancel.clone(); + + tokio::spawn(async move { + loop { + let event = tokio::select! { + _ = cancel.cancelled() => return, + event = events.recv() => event, + }; + let Some(inner) = weak.upgrade() else { + return; + }; + match event { + Ok(ClaimEvent::Delete(key)) => Self::evict_key(&inner, &key), + Ok(ClaimEvent::Reset) | Err(broadcast::error::RecvError::Lagged(_)) => { + Self::clear_entries(&inner); + } + Err(broadcast::error::RecvError::Closed) => { + Self::clear_entries(&inner); + return; + } + } + } + }); + } + + fn evict_key(inner: &AffinityCoordinatorInner, key: &str) { + let Some((_, entry)) = inner.entries.remove(key) else { + return; + }; + if let AffinityEntry::Initializing { notify, .. } = entry { + notify.notify_waiters(); + } + Self::decrement_entry_count(inner, 1); + tracing::debug!(claim_key = key, "evicted session affinity cache entry"); + } + + fn clear_entries(inner: &AffinityCoordinatorInner) { + for entry in inner.entries.iter() { + if let AffinityEntry::Initializing { notify, .. } = entry.value() { + notify.notify_waiters(); + } + } + inner.entries.clear(); + inner.entry_count.store(0, Ordering::Relaxed); + tracing::debug!("cleared session affinity cache after claim watcher reset"); + } + + fn decrement_entry_count(inner: &AffinityCoordinatorInner, removed: usize) { + let _ = inner + .entry_count + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| { + Some(count.saturating_sub(removed)) + }); + } + fn spawn_reaper(inner: &Arc) { let weak = Arc::downgrade(inner); let cancel = inner.cancel.clone(); let period = inner.ttl.min(Duration::from_secs(30)); #[cfg(test)] - let reaper_started = inner.reaper_started.clone(); + let reaper_started = inner.probe.reaper_started.clone(); tokio::spawn(async move { #[cfg(test)] reaper_started.notify_one(); @@ -146,54 +308,48 @@ impl AffinityCoordinator { removed += usize::from(!retain); retain }); - inner.entry_count.fetch_sub(removed, Ordering::Relaxed); + Self::decrement_entry_count(&inner, removed); } }); } #[cfg(test)] - pub async fn acquire( + pub(crate) async fn acquire( &self, session_id: &SessionAffinityId, - requested_target: Option, ) -> Result { - self.acquire_inner(session_id, requested_target, None).await + self.acquire_inner(session_id, None).await } - pub async fn acquire_with_context( + pub(crate) async fn acquire_with_context( &self, session_id: &SessionAffinityId, - requested_target: Option, request_context: &dyn AsyncEngineContext, ) -> Result { - self.acquire_inner(session_id, requested_target, Some(request_context)) - .await + self.acquire_inner(session_id, Some(request_context)).await } async fn acquire_inner( &self, session_id: &SessionAffinityId, - requested_target: Option, request_context: Option<&dyn AsyncEngineContext>, ) -> Result { self.validate_session_id(session_id)?; - let session_id = session_id.as_str().to_string(); + let claim_key = self.inner.claims.key(session_id); loop { let now = Instant::now(); - match self.inner.entries.entry(session_id.clone()) { + match self.inner.entries.entry(claim_key.clone()) { Entry::Vacant(entry) => { self.reserve_entry()?; - return Ok(AffinityAcquire::Initialize(entry.insert_initializing( - &self.inner, - session_id, - requested_target, - ))); + return Ok(AffinityAcquire::Initialize( + entry.insert_initializing(&self.inner, claim_key), + )); } Entry::Occupied(mut entry) => match entry.get_mut() { AffinityEntry::Initializing { notify, .. } => { #[cfg(test)] - self.inner.waiter_observed.notify_one(); + self.inner.probe.waiter_observed.notify_one(); let notified = notify.clone().notified_owned(); tokio::pin!(notified); notified.as_mut().enable(); @@ -228,10 +384,9 @@ impl AffinityCoordinator { drop(entry); return Ok(AffinityAcquire::Initialize(AffinityInitialization { coordinator: Arc::downgrade(&self.inner), - session_id, + claim_key, revision, notify, - requested_target, active: true, })); } @@ -241,11 +396,10 @@ impl AffinityCoordinator { active_leases, .. } => { - validate_bound_target(&session_id, *target, requested_target)?; *active_leases += 1; let lease = AffinityLease { coordinator: Arc::downgrade(&self.inner), - session_id, + claim_key, revision: *revision, active: true, }; @@ -262,10 +416,10 @@ impl AffinityCoordinator { pub fn query_target( &self, session_id: &SessionAffinityId, - requested_target: Option, ) -> Result, Error> { self.validate_session_id(session_id)?; - let Some(entry) = self.inner.entries.get(session_id.as_str()) else { + let claim_key = self.inner.claims.key(session_id); + let Some(entry) = self.inner.entries.get(&claim_key) else { return Ok(None); }; let AffinityEntry::Bound { @@ -280,7 +434,6 @@ impl AffinityCoordinator { if *active_leases == 0 && *idle_deadline <= Instant::now() { return Ok(None); } - validate_bound_target(session_id.as_str(), *target, requested_target)?; Ok(Some(*target)) } @@ -289,6 +442,11 @@ impl AffinityCoordinator { self.inner.entry_count.load(Ordering::Relaxed) } + #[cfg(test)] + pub(super) fn claim_key_for_test(&self, session_id: &SessionAffinityId) -> String { + self.inner.claims.key(session_id) + } + #[cfg(test)] pub(super) fn cancellation_token(&self) -> CancellationToken { self.inner.cancel.clone() @@ -296,17 +454,18 @@ impl AffinityCoordinator { #[cfg(test)] pub(super) async fn wait_for_reaper(&self) { - self.inner.reaper_started.notified().await; + self.inner.probe.reaper_started.notified().await; } #[cfg(test)] pub(super) async fn wait_for_initializing_waiter(&self) { - self.inner.waiter_observed.notified().await; + self.inner.probe.waiter_observed.notified().await; } #[cfg(test)] pub(super) fn expire_for_test(&self, session_id: &SessionAffinityId) { - let Some(mut entry) = self.inner.entries.get_mut(session_id.as_str()) else { + let claim_key = self.inner.claims.key(session_id); + let Some(mut entry) = self.inner.entries.get_mut(&claim_key) else { panic!("session affinity entry missing"); }; let AffinityEntry::Bound { @@ -323,7 +482,14 @@ impl AffinityCoordinator { #[cfg(test)] pub(super) fn with_test_limits(max_entries: usize, max_session_id_bytes: usize) -> Self { - Self::new_with_limits(Duration::from_secs(10), max_entries, max_session_id_bytes).unwrap() + Self::new_with_limits( + Duration::from_secs(10), + max_entries, + max_session_id_bytes, + "local".to_string(), + None, + ) + .unwrap() } fn validate_session_id(&self, session_id: &SessionAffinityId) -> Result<(), Error> { @@ -351,8 +517,7 @@ trait VacantEntryExt { fn insert_initializing( self, inner: &Arc, - session_id: String, - requested_target: Option, + claim_key: String, ) -> AffinityInitialization; } @@ -360,8 +525,7 @@ impl<'a> VacantEntryExt for dashmap::mapref::entry::VacantEntry<'a, String, Affi fn insert_initializing( self, inner: &Arc, - session_id: String, - requested_target: Option, + claim_key: String, ) -> AffinityInitialization { let revision = inner.next_revision.fetch_add(1, Ordering::Relaxed); let notify = Arc::new(Notify::new()); @@ -371,16 +535,15 @@ impl<'a> VacantEntryExt for dashmap::mapref::entry::VacantEntry<'a, String, Affi }); AffinityInitialization { coordinator: Arc::downgrade(inner), - session_id, + claim_key, revision, notify, - requested_target, active: true, } } } -pub enum AffinityAcquire { +pub(crate) enum AffinityAcquire { Initialize(AffinityInitialization), Bound { target: AffinityTarget, @@ -389,56 +552,57 @@ pub enum AffinityAcquire { } impl AffinityAcquire { - pub fn target(&self) -> Option { - match self { - Self::Initialize(_) => None, - Self::Bound { target, .. } => Some(*target), - } - } - - pub fn into_stream( + pub(crate) async fn resolve( self, - selected_target: AffinityTarget, - stream: ManyOut, - ) -> Result, Error> { + proposed_payload: ClaimPayloadFuture<'_>, + ) -> Result { match self { - Self::Initialize(initialization) => { - Ok(initialization.commit(selected_target)?.into_stream(stream)) - } - Self::Bound { target, mut lease } => { - if let Err(error) = validate_bound_target("session", target, Some(selected_target)) - { - lease.invalidate(); - return Err(error); - } - Ok(lease.into_stream(stream)) - } - } - } - - pub fn invalidate(self) { - if let Self::Bound { mut lease, .. } = self { - lease.invalidate(); + Self::Initialize(initialization) => initialization.resolve(proposed_payload).await, + Self::Bound { target, lease } => Ok(ResolvedAffinity { + target, + lease, + created: false, + }), } } } -pub struct AffinityInitialization { +pub(crate) struct AffinityInitialization { coordinator: Weak, - session_id: String, + claim_key: String, revision: u64, notify: Arc, - requested_target: Option, active: bool, } impl AffinityInitialization { - pub fn commit(mut self, target: AffinityTarget) -> Result { - validate_bound_target(&self.session_id, target, self.requested_target)?; + async fn resolve( + self, + mut proposed_payload: ClaimPayloadFuture<'_>, + ) -> Result { + let Some(inner) = self.coordinator.upgrade() else { + return Err(anyhow::anyhow!("session affinity coordinator dropped")); + }; + + let (payload, created) = inner + .claims + .resolve(&self.claim_key, &mut proposed_payload) + .await?; + let target: AffinityTarget = serde_json::from_value(payload) + .map_err(|err| anyhow::anyhow!("invalid session affinity claim payload: {err}"))?; + let lease = self.commit(target)?; + Ok(ResolvedAffinity { + target, + lease, + created, + }) + } + + fn commit(mut self, target: AffinityTarget) -> Result { let Some(inner) = self.coordinator.upgrade() else { return Err(anyhow::anyhow!("session affinity coordinator dropped")); }; - let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { + let Some(mut entry) = inner.entries.get_mut(&self.claim_key) else { return Err(invalid_argument( "session affinity initialization was cancelled", )); @@ -460,7 +624,7 @@ impl AffinityInitialization { self.notify.notify_waiters(); Ok(AffinityLease { coordinator: Arc::downgrade(&inner), - session_id: self.session_id.clone(), + claim_key: self.claim_key.clone(), revision: self.revision, active: true, }) @@ -475,38 +639,27 @@ impl Drop for AffinityInitialization { let Some(inner) = self.coordinator.upgrade() else { return; }; - let removed = inner.entries.remove_if(&self.session_id, |_, entry| { + let removed = inner.entries.remove_if(&self.claim_key, |_, entry| { matches!( entry, AffinityEntry::Initializing { revision, .. } if *revision == self.revision ) }); if removed.is_some() { - inner.entry_count.fetch_sub(1, Ordering::Relaxed); + AffinityCoordinator::decrement_entry_count(&inner, 1); } self.notify.notify_waiters(); } } -pub struct AffinityLease { +pub(crate) struct AffinityLease { coordinator: Weak, - session_id: String, + claim_key: String, revision: u64, active: bool, } impl AffinityLease { - pub fn into_stream(self, stream: ManyOut) -> ManyOut { - let context = stream.context(); - ResponseStream::new( - Box::pin(AffinityTrackedStream { - stream, - lease: Some(self), - }), - context, - ) - } - fn release(&mut self) { if !self.active { return; @@ -515,7 +668,7 @@ impl AffinityLease { let Some(inner) = self.coordinator.upgrade() else { return; }; - let Some(mut entry) = inner.entries.get_mut(&self.session_id) else { + let Some(mut entry) = inner.entries.get_mut(&self.claim_key) else { return; }; let AffinityEntry::Bound { @@ -533,25 +686,6 @@ impl AffinityLease { *active_leases -= 1; *idle_deadline = Instant::now() + inner.ttl; } - - fn invalidate(&mut self) { - if !self.active { - return; - } - self.active = false; - let Some(inner) = self.coordinator.upgrade() else { - return; - }; - let removed = inner.entries.remove_if(&self.session_id, |_, entry| { - matches!( - entry, - AffinityEntry::Bound { revision, .. } if *revision == self.revision - ) - }); - if removed.is_some() { - inner.entry_count.fetch_sub(1, Ordering::Relaxed); - } - } } impl Drop for AffinityLease { @@ -560,9 +694,83 @@ impl Drop for AffinityLease { } } +pub(crate) struct ResolvedAffinity { + target: AffinityTarget, + lease: AffinityLease, + created: bool, +} + +impl ResolvedAffinity { + pub(crate) fn target(&self) -> AffinityTarget { + self.target + } + + pub(crate) fn was_created(&self) -> bool { + self.created + } + + pub(crate) fn into_stream( + self, + stream: ManyOut, + close_on_finish: bool, + ) -> ManyOut { + let context = stream.context(); + let close = close_on_finish.then(|| CloseAction { + coordinator: self.lease.coordinator.clone(), + claims: self + .lease + .coordinator + .upgrade() + .map(|inner| inner.claims.clone()), + claim_key: self.lease.claim_key.clone(), + }); + ResponseStream::new( + Box::pin(AffinityTrackedStream { + stream, + lease: Some(self.lease), + close, + }), + context, + ) + } +} + +struct CloseAction { + coordinator: Weak, + claims: Option, + claim_key: String, +} + +impl CloseAction { + fn run(self) { + if let Some(inner) = self.coordinator.upgrade() { + AffinityCoordinator::evict_key(&inner, &self.claim_key); + } + let Some(claims) = self.claims else { + return; + }; + let claim_key = self.claim_key; + tokio::spawn(async move { + if let Err(error) = claims.close(&claim_key).await { + tracing::error!(%claim_key, %error, "failed to close session affinity claim"); + } + }); + } +} + struct AffinityTrackedStream { stream: ManyOut, lease: Option, + close: Option, +} + +impl AffinityTrackedStream { + fn finish(&mut self) { + drop(self.lease.take()); + if let Some(close) = self.close.take() { + close.run(); + } + } } impl Stream for AffinityTrackedStream { @@ -571,7 +779,7 @@ impl Stream for AffinityTrackedStream { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { match Pin::new(&mut self.stream).poll_next(cx) { Poll::Ready(None) => { - drop(self.lease.take()); + self.finish(); Poll::Ready(None) } Poll::Ready(Some(item)) => Poll::Ready(Some(item)), @@ -580,6 +788,12 @@ impl Stream for AffinityTrackedStream { } } +impl Drop for AffinityTrackedStream { + fn drop(&mut self) { + self.finish(); + } +} + pub fn affinity_id( request: &dynamo_runtime::pipeline::SingleIn, ) -> Result>, Error> { @@ -588,6 +802,13 @@ pub fn affinity_id( .map_err(|message| invalid_argument(format!("invalid session affinity context: {message}"))) } +pub fn session_final(request: &PreprocessedRequest) -> bool { + request + .agent_context + .as_ref() + .is_some_and(|context| context.session_final == Some(true)) +} + pub fn explicit_target( request: &PreprocessedRequest, phase: RequestPhase, @@ -617,31 +838,6 @@ pub fn explicit_target( Ok(worker_id.map(|worker_id| AffinityTarget { worker_id, dp_rank })) } -fn validate_bound_target( - session_id: &str, - bound: AffinityTarget, - requested: Option, -) -> Result<(), Error> { - let Some(requested) = requested else { - return Ok(()); - }; - if bound.worker_id != requested.worker_id { - return Err(invalid_argument(format!( - "session {session_id} is bound to worker {}, not {}", - bound.worker_id, requested.worker_id - ))); - } - match (bound.dp_rank, requested.dp_rank) { - (Some(bound), Some(requested)) if bound != requested => Err(invalid_argument(format!( - "session {session_id} is bound to DP rank {bound}, not {requested}" - ))), - (None, Some(requested)) => Err(invalid_argument(format!( - "session {session_id} has worker-only affinity and cannot add DP rank {requested}" - ))), - _ => Ok(()), - } -} - pub(crate) fn invalid_argument(message: impl Into) -> Error { DynamoError::builder() .error_type(ErrorType::InvalidArgument) diff --git a/lib/llm/src/session_affinity/mod.rs b/lib/llm/src/session_affinity/mod.rs index 08e6451d5533..1f841dab0a6b 100644 --- a/lib/llm/src/session_affinity/mod.rs +++ b/lib/llm/src/session_affinity/mod.rs @@ -4,11 +4,8 @@ mod coordinator; mod push_router; -pub(crate) use coordinator::affinity_id; -pub use coordinator::{ - AffinityAcquire, AffinityCoordinator, AffinityInitialization, AffinityLease, AffinityTarget, - explicit_target, -}; +pub use coordinator::{AffinityCoordinator, AffinityTarget, explicit_target}; +pub(crate) use coordinator::{ResolvedAffinity, affinity_id, session_final}; pub use push_router::SessionAffinityPushRouter; pub const MAX_SESSION_AFFINITY_TTL_SECS: u64 = 31_536_000; diff --git a/lib/llm/src/session_affinity/push_router.rs b/lib/llm/src/session_affinity/push_router.rs index 48f4f6f0b8a5..4d144187691d 100644 --- a/lib/llm/src/session_affinity/push_router.rs +++ b/lib/llm/src/session_affinity/push_router.rs @@ -3,15 +3,19 @@ use std::time::Duration; -use dynamo_runtime::pipeline::{ - AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, Error, ManyOut, PushRouter, - SingleIn, async_trait as pipeline_async_trait, +use dynamo_runtime::{ + discovery::ClaimPayloadFuture, + pipeline::{ + AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, Error, ManyOut, PushRouter, + SingleIn, async_trait as pipeline_async_trait, + }, + traits::DistributedRuntimeProvider, }; use super::{ - AffinityCoordinator, AffinityTarget, LlmResponse, + AffinityCoordinator, AffinityTarget, LlmResponse, ResolvedAffinity, coordinator::{affinity_id, invalid_argument}, - explicit_target, + explicit_target, session_final, }; use crate::{ preprocessor::PreprocessedRequest, @@ -30,9 +34,18 @@ impl SessionAffinityPushRouter { ttl: Option, direct: bool, ) -> Result { + let affinity = ttl + .map(|ttl| { + AffinityCoordinator::new_distributed( + ttl, + inner.client.endpoint.id().to_string(), + inner.client.endpoint.drt().discovery(), + ) + }) + .transpose()?; Ok(Self { inner, - affinity: ttl.map(AffinityCoordinator::new).transpose()?, + affinity, direct, }) } @@ -57,54 +70,44 @@ impl SessionAffinityPushRouter { tracker.record_worker(target.worker_id, target.dp_rank, worker_type); } - fn direct_target( - &self, - explicit: Option, - phase: RequestPhase, - ) -> Result, Error> { - if !self.direct { - return Ok(explicit); - } - explicit.map(Some).ok_or_else(|| { - invalid_argument(format!( - "worker ID required for {phase} request in Direct routing mode" - )) - }) - } - pub fn peek_next_worker(&self) -> Option { self.inner.peek_next_worker() } - async fn acquire_routable( + async fn resolve_affinity( &self, session_id: &crate::protocols::common::extensions::SessionAffinityId, explicit: Option, + request: &PreprocessedRequest, request_context: &dyn AsyncEngineContext, - ) -> Result { + ) -> Result { let affinity = self .affinity .as_ref() .expect("affinity acquisition requires an enabled coordinator"); let operation = affinity - .acquire_with_context(session_id, explicit, request_context) + .acquire_with_context(session_id, request_context) .await?; - let Some(target) = operation.target() else { - return Ok(operation); - }; - if self - .inner - .client - .instance_ids_avail() - .contains(&target.worker_id) - { - return Ok(operation); - } - - operation.invalidate(); - affinity - .acquire_with_context(session_id, explicit, request_context) - .await + let proposed_payload: ClaimPayloadFuture<'_> = Box::pin(async move { + let target = explicit + .or_else(|| { + self.inner + .peek_worker_for_request(request) + .map(|worker_id| AffinityTarget { + worker_id, + dp_rank: None, + }) + }) + .ok_or_else(|| { + if self.direct { + invalid_argument("worker ID required to create Direct session affinity") + } else { + anyhow::anyhow!("no worker available for session affinity") + } + })?; + Ok(serde_json::to_value(target)?) + }); + operation.resolve(proposed_payload).await } pub async fn select_and_dispatch_prefill( @@ -129,14 +132,11 @@ impl SessionAffinityPushRouter { }) .await; } - let explicit = self.direct_target( - explicit_target(&request, RequestPhase::Prefill)?, - RequestPhase::Prefill, - )?; + let explicit = explicit_target(&request, RequestPhase::Prefill)?; let Some(session_id) = session_id else { let Some(pinned_worker) = explicit else { return Err(invalid_argument( - "Direct routing requires an explicit prefill target", + "worker ID required for prefill request in Direct routing mode", )); }; return self @@ -154,7 +154,7 @@ impl SessionAffinityPushRouter { .affinity .as_ref() .expect("affinity query requires an enabled coordinator") - .query_target(&session_id, explicit)? + .query_target(&session_id)? .or(explicit); let rank = selected.and_then(|target| target.dp_rank); return self @@ -174,17 +174,25 @@ impl SessionAffinityPushRouter { .await; } + let close_on_finish = session_final(request.content()); let request_context = request.context(); - let operation = self - .acquire_routable(&session_id, explicit, request_context.as_ref()) + let resolved = self + .resolve_affinity( + &session_id, + explicit, + request.content(), + request_context.as_ref(), + ) .await?; - let selected = operation.target().or(explicit); - let rank = selected.and_then(|target| target.dp_rank); - let dispatch = self + let target = resolved.target(); + let rank = target.dp_rank; + let advance_round_robin = resolved.was_created() && explicit.is_none(); + let ((metadata, _), stream) = self .inner - .select_and_dispatch_exact( + .book_and_dispatch_exact( request, - selected.map(|target| target.worker_id), + target.worker_id, + advance_round_robin, move |request, worker_id| { let target = AffinityTarget { worker_id, @@ -194,15 +202,8 @@ impl SessionAffinityPushRouter { Ok((prepare(request, worker_id, rank)?, target)) }, ) - .await; - let ((metadata, target), stream) = match dispatch { - Ok(result) => result, - Err(error) => { - operation.invalidate(); - return Err(error); - } - }; - Ok((metadata, operation.into_stream(target, stream)?)) + .await?; + Ok((metadata, resolved.into_stream(stream, close_on_finish))) } } @@ -223,11 +224,11 @@ impl AsyncEngine, ManyOut, Error> if !self.direct && session_id.is_none() { return self.inner.generate(request).await; } - let explicit = self.direct_target(explicit_target(&request, phase)?, phase)?; + let explicit = explicit_target(&request, phase)?; let Some(session_id) = session_id else { let Some(target) = explicit else { return Err(invalid_argument(format!( - "Direct routing requires an explicit {phase} target" + "worker ID required for {phase} request in Direct routing mode" ))); }; return self.inner.direct(request, target.worker_id).await; @@ -239,7 +240,7 @@ impl AsyncEngine, ManyOut, Error> .affinity .as_ref() .expect("affinity query requires an enabled coordinator") - .query_target(&session_id, explicit)? + .query_target(&session_id)? .or(explicit); let rank = target.and_then(|target| target.dp_rank); let (_, stream) = self @@ -265,17 +266,25 @@ impl AsyncEngine, ManyOut, Error> return Ok(stream); } + let close_on_finish = session_final(request.content()); let request_context = request.context(); - let operation = self - .acquire_routable(&session_id, explicit, request_context.as_ref()) + let resolved = self + .resolve_affinity( + &session_id, + explicit, + request.content(), + request_context.as_ref(), + ) .await?; - let selected = operation.target().or(explicit); - let rank = selected.and_then(|target| target.dp_rank); - let dispatch = self + let target = resolved.target(); + let rank = target.dp_rank; + let advance_round_robin = resolved.was_created() && explicit.is_none(); + let (_, stream) = self .inner - .select_and_dispatch_exact( + .book_and_dispatch_exact( request, - selected.map(|target| target.worker_id), + target.worker_id, + advance_round_robin, move |request, worker_id| { if rank.is_some() { request.routing_mut().dp_rank = rank; @@ -288,15 +297,8 @@ impl AsyncEngine, ManyOut, Error> Ok(target) }, ) - .await; - let (target, stream) = match dispatch { - Ok(result) => result, - Err(error) => { - operation.invalidate(); - return Err(error); - } - }; - operation.into_stream(target, stream) + .await?; + Ok(resolved.into_stream(stream, close_on_finish)) } } @@ -322,7 +324,6 @@ mod tests { extensions::{SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId}, preprocessor::RoutingHints, }; - use crate::session_affinity::AffinityAcquire; fn request(worker_id: Option, query_only: bool) -> PreprocessedRequest { PreprocessedRequest::builder() @@ -388,51 +389,34 @@ mod tests { } #[tokio::test] - async fn session_affinity_simple_modes_rollback_failed_initialization() { + async fn session_affinity_failed_dispatch_preserves_created_claim() { let runtime = Runtime::from_current().unwrap(); let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) .await .unwrap(); - let namespace = distributed + let client = distributed .namespace("session_affinity_adapters".to_string()) + .unwrap() + .component("workers".to_string()) + .unwrap() + .endpoint("direct") + .client() + .await .unwrap(); - let component = namespace.component("workers".to_string()).unwrap(); - - for (index, mode) in [ - RouterMode::Random, - RouterMode::RoundRobin, - RouterMode::PowerOfTwoChoices, - RouterMode::LeastLoaded, - RouterMode::DeviceAwareWeighted, - RouterMode::Direct, - ] - .into_iter() - .enumerate() - { - let endpoint = component.endpoint(format!("mode-{index}")); - let client = endpoint.client().await.unwrap(); - let inner = PushRouter::from_client(client, mode).await.unwrap(); - let router = SessionAffinityPushRouter::new( - inner, - Some(Duration::from_secs(10)), - mode.is_direct_routing(), - ) + let inner = PushRouter::from_client(client, RouterMode::Direct) + .await .unwrap(); - let worker_id = mode.is_direct_routing().then_some(99); - - assert!( - router - .generate(affinity_request(worker_id, false)) - .await - .is_err() - ); - assert_eq!( - affinity(&router).entry_count(), - 0, - "failed {mode:?} dispatch must release initialization" - ); - } + let router = + SessionAffinityPushRouter::new(inner, Some(Duration::from_secs(10)), true).unwrap(); + + assert!( + router + .generate(affinity_request(Some(99), false)) + .await + .is_err() + ); + assert_eq!(affinity(&router).entry_count(), 1); runtime.shutdown(); } @@ -517,7 +501,7 @@ mod tests { } #[tokio::test] - async fn session_affinity_unavailable_target_is_invalidated() { + async fn session_affinity_unavailable_target_remains_bound() { let runtime = Runtime::from_current().unwrap(); let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) @@ -537,19 +521,20 @@ mod tests { let router = SessionAffinityPushRouter::new(inner, Some(Duration::from_secs(10)), false).unwrap(); let session_id = SessionAffinityId::new("adapter-session"); - let AffinityAcquire::Initialize(initializer) = - affinity(&router).acquire(&session_id, None).await.unwrap() - else { - panic!("first request must initialize"); + let unavailable_target = AffinityTarget { + worker_id: 99, + dp_rank: None, }; - drop( - initializer - .commit(AffinityTarget { - worker_id: 99, - dp_rank: None, - }) - .unwrap(), - ); + let resolved = affinity(&router) + .acquire(&session_id) + .await + .unwrap() + .resolve(Box::pin(async move { + Ok(serde_json::to_value(unavailable_target)?) + })) + .await + .unwrap(); + drop(resolved); assert!( router @@ -558,8 +543,11 @@ mod tests { .is_err() ); assert_eq!( - affinity(&router).query_target(&session_id, None).unwrap(), - None + affinity(&router).query_target(&session_id).unwrap(), + Some(AffinityTarget { + worker_id: 99, + dp_rank: None, + }) ); runtime.shutdown(); diff --git a/lib/llm/src/session_affinity/tests.rs b/lib/llm/src/session_affinity/tests.rs index 044eea963356..8bdd1212f240 100644 --- a/lib/llm/src/session_affinity/tests.rs +++ b/lib/llm/src/session_affinity/tests.rs @@ -1,19 +1,32 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::{sync::Arc, time::Duration}; +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, AtomicUsize, Ordering}, + }, + time::Duration, +}; +use async_trait::async_trait; use dynamo_runtime::{ + discovery::{ + ClaimCloseOutcome, ClaimEvent, ClaimOutcome, ClaimPayload, ClaimPayloadFuture, Discovery, + DiscoveryInstance, DiscoveryQuery, DiscoverySpec, DiscoveryStream, + }, engine::AsyncEngineContext, error::ErrorType, pipeline::{Context, ResponseStream, context::Controller}, protocols::maybe_error::MaybeError, }; use futures::{StreamExt, stream}; +use tokio::sync::broadcast; +use tokio_util::sync::CancellationToken; -use super::{ - AffinityAcquire, AffinityCoordinator, AffinityTarget, LlmResponse, affinity_id, explicit_target, -}; +use super::coordinator::AffinityAcquire; +use super::{AffinityCoordinator, AffinityTarget, LlmResponse, affinity_id, explicit_target}; use crate::{ preprocessor::PreprocessedRequest, protocols::common::{ @@ -37,6 +50,138 @@ fn coordinator() -> AffinityCoordinator { AffinityCoordinator::new(Duration::from_secs(10)).unwrap() } +struct ClaimTestDiscovery { + claims: Mutex>, + create_calls: AtomicUsize, + close_calls: AtomicUsize, + events: Mutex>>, +} + +impl ClaimTestDiscovery { + fn new(event_capacity: usize) -> Arc { + let (events, _) = broadcast::channel(event_capacity); + Arc::new(Self { + claims: Mutex::new(HashMap::new()), + create_calls: AtomicUsize::new(0), + close_calls: AtomicUsize::new(0), + events: Mutex::new(Some(events)), + }) + } + + fn emit(&self, event: ClaimEvent) { + if let Some(events) = self.events.lock().unwrap().as_ref() { + let _ = events.send(event); + } + } + + fn disconnect(&self) { + self.events.lock().unwrap().take(); + } +} + +#[async_trait] +impl Discovery for ClaimTestDiscovery { + fn instance_id(&self) -> u64 { + 1 + } + + async fn register_internal(&self, _spec: DiscoverySpec) -> anyhow::Result { + anyhow::bail!("not used by claim tests") + } + + async fn unregister(&self, _instance: DiscoveryInstance) -> anyhow::Result<()> { + Ok(()) + } + + async fn list(&self, _query: DiscoveryQuery) -> anyhow::Result> { + Ok(Vec::new()) + } + + async fn list_and_watch( + &self, + _query: DiscoveryQuery, + _cancel_token: Option, + ) -> anyhow::Result { + Ok(Box::pin(stream::pending())) + } + + async fn create_or_get_claim( + &self, + key: &str, + proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> anyhow::Result { + self.create_calls.fetch_add(1, Ordering::Relaxed); + if let Some(payload) = self.claims.lock().unwrap().get(key).cloned() { + return Ok(ClaimOutcome::Existing(payload)); + } + + let proposed = proposed_payload.as_mut().await?; + let mut claims = self.claims.lock().unwrap(); + if let Some(payload) = claims.get(key).cloned() { + return Ok(ClaimOutcome::Existing(payload)); + } + claims.insert(key.to_string(), proposed.clone()); + Ok(ClaimOutcome::Created(proposed)) + } + + async fn close_claim(&self, key: &str) -> anyhow::Result { + self.close_calls.fetch_add(1, Ordering::Relaxed); + if self.claims.lock().unwrap().remove(key).is_some() { + self.emit(ClaimEvent::Delete(key.to_string())); + } + Ok(ClaimCloseOutcome::Closed) + } + + fn subscribe_claim_events(&self) -> Option> { + self.events + .lock() + .unwrap() + .as_ref() + .map(broadcast::Sender::subscribe) + } +} + +fn distributed_coordinator(discovery: Arc) -> AffinityCoordinator { + AffinityCoordinator::new_distributed( + Duration::from_secs(10), + "ns/component/endpoint".to_string(), + discovery, + ) + .unwrap() +} + +fn target_payload(target: AffinityTarget) -> ClaimPayloadFuture<'static> { + Box::pin(async move { Ok(serde_json::to_value(target)?) }) +} + +async fn resolve_local( + coordinator: &AffinityCoordinator, + session_id: &SessionAffinityId, + selected: AffinityTarget, +) -> super::ResolvedAffinity { + coordinator + .acquire(session_id) + .await + .unwrap() + .resolve(target_payload(selected)) + .await + .unwrap() +} + +async fn wait_for_cached_target( + coordinator: &AffinityCoordinator, + session_id: &SessionAffinityId, + expected: Option, +) { + for _ in 0..100 { + if coordinator.query_target(session_id).unwrap() == expected { + return; + } + tokio::task::yield_now().await; + } + assert_eq!(coordinator.query_target(session_id).unwrap(), expected); +} + fn response_stream(items: usize) -> dynamo_runtime::pipeline::ManyOut { let items = (0..items).map(|_| Annotated::from_data(LLMEngineOutput::default())); ResponseStream::new( @@ -61,11 +206,11 @@ fn cancelled_response_stream() -> dynamo_runtime::pipeline::ManyOut async fn assert_binding_expires_after_refreshed_ttl(coordinator: &AffinityCoordinator) { tokio::time::advance(Duration::from_secs(9)).await; assert_eq!( - coordinator.query_target(&session_id(), None).unwrap(), + coordinator.query_target(&session_id()).unwrap(), Some(target(7, Some(0))) ); tokio::time::advance(Duration::from_secs(2)).await; - assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); + assert_eq!(coordinator.query_target(&session_id()).unwrap(), None); } fn request_with_routing(routing: RoutingHints) -> PreprocessedRequest { @@ -136,40 +281,207 @@ fn session_affinity_context_type_errors_are_preserved() { #[tokio::test(start_paused = true)] async fn session_affinity_initialization_is_atomic() { let coordinator = coordinator(); - let first = coordinator.acquire(&session_id(), None).await.unwrap(); - let AffinityAcquire::Initialize(first) = first else { - panic!("first request must initialize"); - }; + let first = coordinator.acquire(&session_id()).await.unwrap(); + assert!(matches!(&first, AffinityAcquire::Initialize(_))); let waiter_coordinator = coordinator.clone(); - let waiter = tokio::spawn(async move { waiter_coordinator.acquire(&session_id(), None).await }); + let waiter = tokio::spawn(async move { waiter_coordinator.acquire(&session_id()).await }); coordinator.wait_for_initializing_waiter().await; assert!(!waiter.is_finished()); - let first_lease = first.commit(target(7, Some(0))).unwrap(); - let second = waiter.await.unwrap().unwrap(); - let AffinityAcquire::Bound { - target: second_target, - lease: second_lease, - } = second - else { - panic!("waiter must acquire the committed binding"); - }; - assert_eq!(second_target, target(7, Some(0))); - drop(first_lease); - drop(second_lease); + let first = first + .resolve(target_payload(target(7, Some(0)))) + .await + .unwrap(); + let second = waiter + .await + .unwrap() + .unwrap() + .resolve(target_payload(target(8, Some(0)))) + .await + .unwrap(); + assert_eq!(second.target(), target(7, Some(0))); + drop(first); + drop(second); +} + +#[tokio::test] +async fn distributed_existing_winner_skips_proposal_and_cache_hits_skip_discovery() { + let discovery = ClaimTestDiscovery::new(16); + let first = distributed_coordinator(discovery.clone()); + let second = distributed_coordinator(discovery.clone()); + let session_id = session_id(); + + let created = first + .acquire(&session_id) + .await + .unwrap() + .resolve(target_payload(target(7, Some(0)))) + .await + .unwrap(); + assert_eq!(created.target(), target(7, Some(0))); + assert!(created.was_created()); + drop(created); + + let proposal_polled = Arc::new(AtomicBool::new(false)); + let polled = proposal_polled.clone(); + let proposed: ClaimPayloadFuture<'_> = Box::pin(async move { + polled.store(true, Ordering::Relaxed); + Ok(serde_json::to_value(target(8, Some(0)))?) + }); + let winner = second + .acquire(&session_id) + .await + .unwrap() + .resolve(proposed) + .await + .unwrap(); + assert_eq!(winner.target(), target(7, Some(0))); + assert!(!winner.was_created()); + assert!(!proposal_polled.load(Ordering::Relaxed)); + drop(winner); + assert_eq!(discovery.create_calls.load(Ordering::Relaxed), 2); + + let cached_proposal_polled = Arc::new(AtomicBool::new(false)); + let polled = cached_proposal_polled.clone(); + let proposed: ClaimPayloadFuture<'_> = Box::pin(async move { + polled.store(true, Ordering::Relaxed); + Ok(serde_json::to_value(target(9, Some(0)))?) + }); + let cached = second + .acquire(&session_id) + .await + .unwrap() + .resolve(proposed) + .await + .unwrap(); + assert_eq!(cached.target(), target(7, Some(0))); + assert!(!cached_proposal_polled.load(Ordering::Relaxed)); + assert_eq!(discovery.create_calls.load(Ordering::Relaxed), 2); +} + +#[tokio::test] +async fn distributed_delete_evicts_one_entry_and_duplicate_is_harmless() { + let discovery = ClaimTestDiscovery::new(16); + let first = distributed_coordinator(discovery.clone()); + let second = distributed_coordinator(discovery.clone()); + let first_session = SessionAffinityId::new("first-session"); + let second_session = SessionAffinityId::new("second-session"); + + for (coordinator, session, worker_id) in + [(&first, &first_session, 7), (&second, &second_session, 8)] + { + drop( + coordinator + .acquire(session) + .await + .unwrap() + .resolve(target_payload(target(worker_id, Some(0)))) + .await + .unwrap(), + ); + } + + let first_key = first.claim_key_for_test(&first_session); + discovery.emit(ClaimEvent::Delete(first_key.clone())); + wait_for_cached_target(&first, &first_session, None).await; + assert_eq!( + second.query_target(&second_session).unwrap(), + Some(target(8, Some(0))) + ); + discovery.emit(ClaimEvent::Delete(first_key)); + wait_for_cached_target(&first, &first_session, None).await; +} + +#[tokio::test] +async fn distributed_subscriber_lag_clears_all_entries() { + let discovery = ClaimTestDiscovery::new(1); + let coordinator = distributed_coordinator(discovery.clone()); + let session_id = session_id(); + drop( + coordinator + .acquire(&session_id) + .await + .unwrap() + .resolve(target_payload(target(8, Some(0)))) + .await + .unwrap(), + ); + + for index in 0..16 { + discovery.emit(ClaimEvent::Delete(format!("unrelated/{index}"))); + } + wait_for_cached_target(&coordinator, &session_id, None).await; +} + +#[tokio::test] +async fn distributed_reset_and_disconnect_clear_entries() { + let discovery = ClaimTestDiscovery::new(16); + let coordinator = distributed_coordinator(discovery.clone()); + let session_id = session_id(); + + drop( + coordinator + .acquire(&session_id) + .await + .unwrap() + .resolve(target_payload(target(8, Some(0)))) + .await + .unwrap(), + ); + discovery.emit(ClaimEvent::Reset); + wait_for_cached_target(&coordinator, &session_id, None).await; + + drop( + coordinator + .acquire(&session_id) + .await + .unwrap() + .resolve(target_payload(target(99, Some(0)))) + .await + .unwrap(), + ); + discovery.disconnect(); + wait_for_cached_target(&coordinator, &session_id, None).await; +} + +#[tokio::test] +async fn terminal_close_evicts_synchronously() { + let discovery = ClaimTestDiscovery::new(16); + let coordinator = distributed_coordinator(discovery.clone()); + let session_id = session_id(); + let resolved = coordinator + .acquire(&session_id) + .await + .unwrap() + .resolve(target_payload(target(7, Some(0)))) + .await + .unwrap(); + let key = coordinator.claim_key_for_test(&session_id); + let mut stream = resolved.into_stream(response_stream(0), true); + assert!(stream.next().await.is_none()); + assert_eq!(coordinator.query_target(&session_id).unwrap(), None); + + for _ in 0..100 { + if discovery.close_calls.load(Ordering::Relaxed) == 1 { + break; + } + tokio::task::yield_now().await; + } + assert_eq!(discovery.close_calls.load(Ordering::Relaxed), 1); + assert!(!discovery.claims.lock().unwrap().contains_key(&key)); } #[tokio::test(start_paused = true)] async fn session_affinity_initializer_cancellation_wakes_waiter() { let coordinator = coordinator(); - let first = coordinator.acquire(&session_id(), None).await.unwrap(); + let first = coordinator.acquire(&session_id()).await.unwrap(); let AffinityAcquire::Initialize(first) = first else { panic!("first request must initialize"); }; let waiter_coordinator = coordinator.clone(); - let waiter = tokio::spawn(async move { waiter_coordinator.acquire(&session_id(), None).await }); + let waiter = tokio::spawn(async move { waiter_coordinator.acquire(&session_id()).await }); coordinator.wait_for_initializing_waiter().await; drop(first); @@ -178,7 +490,7 @@ async fn session_affinity_initializer_cancellation_wakes_waiter() { drop(next); assert_eq!(coordinator.entry_count(), 0); assert!(matches!( - coordinator.acquire(&session_id(), None).await.unwrap(), + coordinator.acquire(&session_id()).await.unwrap(), AffinityAcquire::Initialize(_) )); } @@ -186,7 +498,7 @@ async fn session_affinity_initializer_cancellation_wakes_waiter() { #[tokio::test(start_paused = true)] async fn session_affinity_wait_stops_when_request_is_cancelled() { let coordinator = coordinator(); - let first = coordinator.acquire(&session_id(), None).await.unwrap(); + let first = coordinator.acquire(&session_id()).await.unwrap(); let AffinityAcquire::Initialize(first) = first else { panic!("first request must initialize"); }; @@ -196,7 +508,7 @@ async fn session_affinity_wait_stops_when_request_is_cancelled() { let waiter_coordinator = coordinator.clone(); let waiter = tokio::spawn(async move { waiter_coordinator - .acquire_with_context(&session_id(), None, waiter_context.as_ref()) + .acquire_with_context(&session_id(), waiter_context.as_ref()) .await }); coordinator.wait_for_initializing_waiter().await; @@ -214,70 +526,28 @@ async fn session_affinity_wait_stops_when_request_is_cancelled() { } #[tokio::test(start_paused = true)] -async fn session_affinity_validates_worker_and_rank_contract() { +async fn session_affinity_existing_binding_overrides_explicit_proposals() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = coordinator - .acquire(&session_id(), Some(target(7, None))) - .await - .unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, None)).unwrap()); + drop(resolve_local(&coordinator, &session_id(), target(7, None)).await); - assert!( - coordinator - .acquire(&session_id(), Some(target(8, None))) + for proposal in [target(8, None), target(7, Some(0)), target(7, None)] { + let resolved = coordinator + .acquire(&session_id()) .await - .is_err() - ); - assert!( - coordinator - .acquire(&session_id(), Some(target(7, Some(0)))) - .await - .is_err() - ); - assert!( - coordinator - .acquire(&session_id(), Some(target(7, None))) + .unwrap() + .resolve(target_payload(proposal)) .await - .is_ok() - ); -} - -#[tokio::test(start_paused = true)] -async fn session_affinity_failed_bound_operation_invalidates_binding() { - let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, Some(0))).unwrap()); - - let operation = coordinator.acquire(&session_id(), None).await.unwrap(); - assert_eq!(operation.target(), Some(target(7, Some(0)))); - operation.invalidate(); - - assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); - assert_eq!(coordinator.entry_count(), 0); - assert!(matches!( - coordinator.acquire(&session_id(), None).await.unwrap(), - AffinityAcquire::Initialize(_) - )); + .unwrap(); + assert_eq!(resolved.target(), target(7, None)); + } } #[tokio::test(start_paused = true)] async fn session_affinity_stream_drop_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let resolved = resolve_local(&coordinator, &session_id(), target(7, Some(0))).await; tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(response_stream(1)); + let mut stream = resolved.into_stream(response_stream(1), false); assert!(stream.next().await.is_some()); drop(stream); @@ -287,14 +557,9 @@ async fn session_affinity_stream_drop_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_empty_stream_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let resolved = resolve_local(&coordinator, &session_id(), target(7, Some(0))).await; tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(response_stream(0)); + let mut stream = resolved.into_stream(response_stream(0), false); assert!(stream.next().await.is_none()); assert_binding_expires_after_refreshed_ttl(&coordinator).await; @@ -303,23 +568,12 @@ async fn session_affinity_empty_stream_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_cancelled_stream_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(resolve_local(&coordinator, &session_id(), target(7, Some(0))).await); tokio::time::advance(Duration::from_secs(9)).await; - let AffinityAcquire::Bound { - target: bound_target, - lease, - } = coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("continuation must acquire the existing binding"); - }; - assert_eq!(bound_target, target(7, Some(0))); - let mut stream = lease.into_stream(cancelled_response_stream()); + let resolved = resolve_local(&coordinator, &session_id(), target(8, Some(0))).await; + assert_eq!(resolved.target(), target(7, Some(0))); + let mut stream = resolved.into_stream(cancelled_response_stream(), false); assert!(stream.next().await.is_none()); assert_binding_expires_after_refreshed_ttl(&coordinator).await; @@ -328,10 +582,8 @@ async fn session_affinity_cancelled_stream_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_committed_binding_survives_cancelled_stream_until_ttl() { let coordinator = coordinator(); - let operation = coordinator.acquire(&session_id(), None).await.unwrap(); - let mut stream = operation - .into_stream(target(7, Some(0)), cancelled_response_stream()) - .unwrap(); + let resolved = resolve_local(&coordinator, &session_id(), target(7, Some(0))).await; + let mut stream = resolved.into_stream(cancelled_response_stream(), false); tokio::time::advance(Duration::from_secs(9)).await; assert!(stream.next().await.is_none()); @@ -341,14 +593,9 @@ async fn session_affinity_committed_binding_survives_cancelled_stream_until_ttl( #[tokio::test(start_paused = true)] async fn session_affinity_error_stream_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let resolved = resolve_local(&coordinator, &session_id(), target(7, Some(0))).await; tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(error_response_stream()); + let mut stream = resolved.into_stream(error_response_stream(), false); assert!(stream.next().await.unwrap().is_err()); assert!(stream.next().await.is_none()); @@ -358,14 +605,9 @@ async fn session_affinity_error_stream_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_stream_eof_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - let lease = initializer.commit(target(7, Some(0))).unwrap(); + let resolved = resolve_local(&coordinator, &session_id(), target(7, Some(0))).await; tokio::time::advance(Duration::from_secs(9)).await; - let mut stream = lease.into_stream(response_stream(1)); + let mut stream = resolved.into_stream(response_stream(1), false); while stream.next().await.is_some() {} assert_binding_expires_after_refreshed_ttl(&coordinator).await; @@ -374,26 +616,17 @@ async fn session_affinity_stream_eof_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_bound_lease_drop_refreshes_idle_ttl() { let coordinator = coordinator(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(resolve_local(&coordinator, &session_id(), target(7, Some(0))).await); tokio::time::advance(Duration::from_secs(9)).await; - let AffinityAcquire::Bound { lease, .. } = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("continuation must acquire the binding"); - }; + let resolved = resolve_local(&coordinator, &session_id(), target(8, Some(0))).await; tokio::time::advance(Duration::from_secs(2)).await; tokio::task::yield_now().await; assert_eq!( - coordinator.query_target(&session_id(), None).unwrap(), + coordinator.query_target(&session_id()).unwrap(), Some(target(7, Some(0))) ); - drop(lease); + drop(resolved); assert_binding_expires_after_refreshed_ttl(&coordinator).await; } @@ -401,27 +634,22 @@ async fn session_affinity_bound_lease_drop_refreshes_idle_ttl() { #[tokio::test(start_paused = true)] async fn session_affinity_query_is_read_only() { let coordinator = coordinator(); - assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); + assert_eq!(coordinator.query_target(&session_id()).unwrap(), None); assert_eq!(coordinator.entry_count(), 0); - let initializing = coordinator.acquire(&session_id(), None).await.unwrap(); - assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); + let initializing = coordinator.acquire(&session_id()).await.unwrap(); + assert_eq!(coordinator.query_target(&session_id()).unwrap(), None); assert_eq!(coordinator.entry_count(), 1); drop(initializing); assert_eq!(coordinator.entry_count(), 0); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(resolve_local(&coordinator, &session_id(), target(7, Some(0))).await); assert_eq!( - coordinator.query_target(&session_id(), None).unwrap(), + coordinator.query_target(&session_id()).unwrap(), Some(target(7, Some(0))) ); coordinator.expire_for_test(&session_id()); - assert_eq!(coordinator.query_target(&session_id(), None).unwrap(), None); + assert_eq!(coordinator.query_target(&session_id()).unwrap(), None); assert_eq!(coordinator.entry_count(), 1); } @@ -429,12 +657,7 @@ async fn session_affinity_query_is_read_only() { async fn session_affinity_reaper_removes_idle_entries_and_stops_on_drop() { let coordinator = coordinator(); let cancellation = coordinator.cancellation_token(); - let AffinityAcquire::Initialize(initializer) = - coordinator.acquire(&session_id(), None).await.unwrap() - else { - panic!("first request must initialize"); - }; - drop(initializer.commit(target(7, Some(0))).unwrap()); + drop(resolve_local(&coordinator, &session_id(), target(7, Some(0))).await); coordinator.wait_for_reaper().await; tokio::time::advance(Duration::from_secs(10)).await; @@ -467,7 +690,7 @@ fn session_affinity_rejects_invalid_ttl_before_starting_reaper() { async fn session_affinity_enforces_id_and_entry_limits() { let coordinator = AffinityCoordinator::with_test_limits(1, 8); let oversized = SessionAffinityId::new("123456789"); - let Err(error) = coordinator.acquire(&oversized, None).await else { + let Err(error) = coordinator.acquire(&oversized).await else { panic!("oversized session ID must fail"); }; assert!(dynamo_runtime::error::match_error_chain( @@ -478,9 +701,9 @@ async fn session_affinity_enforces_id_and_entry_limits() { assert_eq!(coordinator.entry_count(), 0); let first_id = SessionAffinityId::new("first"); - let first = coordinator.acquire(&first_id, None).await.unwrap(); + let first = coordinator.acquire(&first_id).await.unwrap(); let second_id = SessionAffinityId::new("second"); - let Err(error) = coordinator.acquire(&second_id, None).await else { + let Err(error) = coordinator.acquire(&second_id).await else { panic!("entry limit must reject a second session"); }; assert!(dynamo_runtime::error::match_error_chain( @@ -492,7 +715,7 @@ async fn session_affinity_enforces_id_and_entry_limits() { drop(first); assert_eq!(coordinator.entry_count(), 0); assert!(matches!( - coordinator.acquire(&second_id, None).await.unwrap(), + coordinator.acquire(&second_id).await.unwrap(), AffinityAcquire::Initialize(_) )); } diff --git a/lib/runtime/src/discovery/kube.rs b/lib/runtime/src/discovery/kube.rs index a07f0bee2137..611f2c2862c1 100644 --- a/lib/runtime/src/discovery/kube.rs +++ b/lib/runtime/src/discovery/kube.rs @@ -15,14 +15,16 @@ use utils::{KubeDiscoveryMode, PodInfo}; use crate::CancellationToken; use crate::discovery::{ - Discovery, DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryMetadata, - DiscoveryQuery, DiscoverySpec, DiscoveryStream, MetadataSnapshot, + ClaimCloseOutcome, ClaimOutcome, ClaimPayloadFuture, Discovery, DiscoveryEvent, + DiscoveryInstance, DiscoveryInstanceId, DiscoveryMetadata, DiscoveryQuery, DiscoverySpec, + DiscoveryStream, MetadataSnapshot, }; use anyhow::Result; use async_trait::async_trait; use kube::{Api, Client as KubeClient, api::DeleteParams}; use std::collections::HashSet; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use tokio::sync::RwLock; /// Kubernetes-based discovery client @@ -33,6 +35,7 @@ pub struct KubeDiscoveryClient { metadata_watch: tokio::sync::watch::Receiver>, kube_client: KubeClient, pod_info: PodInfo, + claim_warning_emitted: Arc, } impl KubeDiscoveryClient { @@ -104,8 +107,19 @@ impl KubeDiscoveryClient { metadata_watch: watch_rx, kube_client, pod_info, + claim_warning_emitted: Arc::new(AtomicBool::new(false)), }) } + + fn warn_claims_unsupported(&self) { + if self.claim_warning_emitted.swap(true, Ordering::Relaxed) { + return; + } + + tracing::warn!( + "Kubernetes discovery does not coordinate session affinity across frontend processes; using process-local affinity" + ); + } } #[async_trait] @@ -496,4 +510,18 @@ impl Discovery for KubeDiscoveryClient { let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(event_rx); Ok(Box::pin(stream)) } + + async fn create_or_get_claim( + &self, + _key: &str, + _proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> Result { + self.warn_claims_unsupported(); + Ok(ClaimOutcome::Unsupported) + } + + async fn close_claim(&self, _key: &str) -> Result { + self.warn_claims_unsupported(); + Ok(ClaimCloseOutcome::Unsupported) + } } diff --git a/lib/runtime/src/discovery/kv_store.rs b/lib/runtime/src/discovery/kv_store.rs index c72e465f3283..26f0fbec60e2 100644 --- a/lib/runtime/src/discovery/kv_store.rs +++ b/lib/runtime/src/discovery/kv_store.rs @@ -3,27 +3,123 @@ use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; -use anyhow::Result; +use anyhow::{Context, Result}; use async_trait::async_trait; use futures::{Stream, StreamExt}; +use tokio::sync::{OnceCell, broadcast, oneshot}; use tokio_util::sync::CancellationToken; use super::{ - Discovery, DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery, - DiscoverySpec, DiscoveryStream, EndpointInstanceId, EventChannelInstanceId, - ModelCardInstanceId, + ClaimCloseOutcome, ClaimEvent, ClaimOutcome, ClaimPayload, ClaimPayloadFuture, Discovery, + DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery, DiscoverySpec, + DiscoveryStream, EndpointInstanceId, EventChannelInstanceId, ModelCardInstanceId, }; use crate::storage::kv; const INSTANCES_BUCKET: &str = "v1/instances"; const MODELS_BUCKET: &str = "v1/mdc"; const EVENT_CHANNELS_BUCKET: &str = "v1/event_channels"; +const CLAIMS_BUCKET: &str = "v1/claims"; +const CLAIM_CREATE_ATTEMPTS: usize = 3; +const CLAIM_WATCH_RECONNECT_BACKOFF: Duration = Duration::from_millis(250); /// Discovery implementation backed by a kv::Store pub struct KVStoreDiscovery { store: Arc, cancel_token: CancellationToken, + claims: ClaimState, +} + +/// Process-local invalidation relay for the shared claims bucket. +/// +/// One backend watcher serves all affinity coordinators attached to this discovery +/// instance. `Put` events never populate caches. `Delete(key)` evicts one entry, while +/// watcher loss, restart, or subscriber lag produces `Reset` so coordinators clear all +/// entries rather than retain potentially stale bindings. +struct ClaimState { + events: broadcast::Sender, + watcher_started: OnceCell<()>, + memory_warning_emitted: AtomicBool, + #[cfg(test)] + watcher_probe: ClaimWatcherProbe, +} + +impl ClaimState { + fn new() -> Self { + let (events, _) = broadcast::channel(1024); + Self { + events, + watcher_started: OnceCell::new(), + memory_warning_emitted: AtomicBool::new(false), + #[cfg(test)] + watcher_probe: ClaimWatcherProbe::new(), + } + } + + fn subscribe(&self) -> broadcast::Receiver { + self.events.subscribe() + } + + fn warn_if_memory(&self, is_memory: bool) { + if !is_memory || self.memory_warning_emitted.swap(true, Ordering::Relaxed) { + return; + } + + tracing::warn!( + "session affinity claims use MemoryStore and coordinate only within this process/store" + ); + } +} + +#[cfg(test)] +struct ClaimWatcherProbe { + start_count: Arc, + active_count: Arc, +} + +#[cfg(test)] +impl ClaimWatcherProbe { + fn new() -> Self { + Self { + start_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + active_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + } + } + + fn record_start(&self) -> Arc { + self.start_count + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.active_count.clone() + } + + fn starts(&self) -> usize { + self.start_count.load(std::sync::atomic::Ordering::Relaxed) + } + + fn active(&self) -> usize { + self.active_count.load(std::sync::atomic::Ordering::Relaxed) + } +} + +#[cfg(test)] +struct ClaimWatcherActiveGuard(Arc); + +#[cfg(test)] +impl ClaimWatcherActiveGuard { + fn new(active_count: Arc) -> Self { + active_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + Self(active_count) + } +} + +#[cfg(test)] +impl Drop for ClaimWatcherActiveGuard { + fn drop(&mut self) { + self.0.fetch_sub(1, std::sync::atomic::Ordering::Relaxed); + } } impl KVStoreDiscovery { @@ -31,7 +127,167 @@ impl KVStoreDiscovery { Self { store: Arc::new(store), cancel_token, + claims: ClaimState::new(), + } + } + + async fn ensure_claim_watcher(&self) -> Result<()> { + self.claims + .watcher_started + .get_or_try_init(|| async { + let (ready_tx, ready_rx) = oneshot::channel(); + let store = self.store.clone(); + let cancel_token = self.cancel_token.clone(); + let claim_events = self.claims.events.clone(); + #[cfg(test)] + let active_count = self.claims.watcher_probe.record_start(); + + tokio::spawn(async move { + #[cfg(test)] + let _active_guard = ClaimWatcherActiveGuard::new(active_count); + Self::run_claim_watcher(store, cancel_token, claim_events, ready_tx).await; + }); + + ready_rx + .await + .context("claim watcher stopped before startup completed")? + .map_err(anyhow::Error::msg) + }) + .await?; + Ok(()) + } + + async fn run_claim_watcher( + store: Arc, + cancel_token: CancellationToken, + claim_events: broadcast::Sender, + ready_tx: oneshot::Sender>, + ) { + let mut ready_tx = Some(ready_tx); + + loop { + if cancel_token.is_cancelled() { + if let Some(ready_tx) = ready_tx.take() { + let _ = ready_tx.send(Err("claim watcher startup was cancelled".to_string())); + } + let _ = claim_events.send(ClaimEvent::Reset); + return; + } + + let bucket = match store.get_or_create_bucket(CLAIMS_BUCKET, None).await { + Ok(bucket) => bucket, + Err(err) => { + if let Some(ready_tx) = ready_tx.take() { + let _ = ready_tx.send(Err(err.to_string())); + return; + } + tracing::error!(error = %err, "failed to reconnect session claim watcher"); + let _ = claim_events.send(ClaimEvent::Reset); + if Self::wait_for_claim_watcher_retry(&cancel_token).await { + return; + } + continue; + } + }; + + let mut stream = match bucket.watch().await { + Ok(stream) => stream, + Err(err) => { + if let Some(ready_tx) = ready_tx.take() { + let _ = ready_tx.send(Err(err.to_string())); + return; + } + tracing::error!(error = %err, "failed to reconnect session claim watch stream"); + let _ = claim_events.send(ClaimEvent::Reset); + if Self::wait_for_claim_watcher_retry(&cancel_token).await { + return; + } + continue; + } + }; + + if let Some(ready_tx) = ready_tx.take() { + let _ = ready_tx.send(Ok(())); + } else { + let _ = claim_events.send(ClaimEvent::Reset); + } + + loop { + let event = tokio::select! { + _ = cancel_token.cancelled() => { + let _ = claim_events.send(ClaimEvent::Reset); + return; + } + event = stream.next() => event, + }; + + let Some(event) = event else { + tracing::warn!( + "session claim watch stream ended; clearing local affinity caches" + ); + let _ = claim_events.send(ClaimEvent::Reset); + break; + }; + + if let kv::WatchEvent::Delete(key) = event { + let key = Self::strip_bucket_prefix(key.as_ref(), CLAIMS_BUCKET).to_string(); + let _ = claim_events.send(ClaimEvent::Delete(key)); + } + } + + if Self::wait_for_claim_watcher_retry(&cancel_token).await { + return; + } + } + } + + async fn wait_for_claim_watcher_retry(cancel_token: &CancellationToken) -> bool { + tokio::select! { + _ = cancel_token.cancelled() => true, + _ = tokio::time::sleep(CLAIM_WATCH_RECONNECT_BACKOFF) => false, + } + } + + fn warn_if_memory_claims(&self) { + self.claims.warn_if_memory(self.store.is_memory()); + } + + fn parse_claim(value: &[u8]) -> Result { + serde_json::from_slice(value).context("failed to deserialize session affinity claim") + } + + async fn create_or_get_in_bucket( + bucket: &dyn kv::Bucket, + key: &kv::Key, + proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> Result { + if let Some(payload) = bucket.get(key).await? { + return Ok(ClaimOutcome::Existing(Self::parse_claim(&payload)?)); } + + let proposed_payload = proposed_payload.as_mut().await?; + let proposed_bytes = serde_json::to_vec(&proposed_payload)?; + + for attempt in 0..CLAIM_CREATE_ATTEMPTS { + match bucket.insert(key, proposed_bytes.clone().into(), 0).await? { + kv::StoreOutcome::Created(_) => { + return Ok(ClaimOutcome::Created(proposed_payload)); + } + kv::StoreOutcome::Exists(_) => { + if let Some(payload) = bucket.get(key).await? { + return Ok(ClaimOutcome::Existing(Self::parse_claim(&payload)?)); + } + + if attempt + 1 == CLAIM_CREATE_ATTEMPTS { + anyhow::bail!( + "session affinity claim disappeared after {CLAIM_CREATE_ATTEMPTS} competing insert attempts" + ); + } + } + } + } + + unreachable!("claim creation loop always returns") } /// Build the key path for an endpoint (relative to bucket, not absolute) @@ -592,6 +848,41 @@ impl Discovery for KVStoreDiscovery { Ok(Box::pin(stream)) } + async fn create_or_get_claim( + &self, + key: &str, + proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> Result { + self.warn_if_memory_claims(); + self.ensure_claim_watcher().await?; + + let bucket = self.store.get_or_create_bucket(CLAIMS_BUCKET, None).await?; + let key = kv::Key::new(key.to_string()); + + Self::create_or_get_in_bucket(bucket.as_ref(), &key, proposed_payload).await + } + + async fn close_claim(&self, key: &str) -> Result { + self.warn_if_memory_claims(); + self.ensure_claim_watcher().await?; + + let Some(bucket) = self.store.get_bucket(CLAIMS_BUCKET).await? else { + return Ok(ClaimCloseOutcome::Closed); + }; + let key = kv::Key::new(key.to_string()); + + match bucket.delete(&key).await { + Ok(()) | Err(kv::StoreError::MissingBucket(_) | kv::StoreError::MissingKey(_)) => { + Ok(ClaimCloseOutcome::Closed) + } + Err(err) => Err(err.into()), + } + } + + fn subscribe_claim_events(&self) -> Option> { + Some(self.claims.subscribe()) + } + fn shutdown(&self) { self.store.shutdown(); } @@ -599,9 +890,220 @@ impl Discovery for KVStoreDiscovery { #[cfg(test)] mod tests { + use std::collections::HashMap; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use super::*; use crate::component::TransportType; + fn payload(worker_id: u64) -> ClaimPayload { + serde_json::json!({"worker_id": worker_id, "dp_rank": 0}) + } + + struct DisappearingBucket { + insert_calls: AtomicUsize, + create_on_call: Option, + } + + #[async_trait] + impl kv::Bucket for DisappearingBucket { + async fn insert( + &self, + _key: &kv::Key, + _value: bytes::Bytes, + _revision: u64, + ) -> std::result::Result { + let call = self.insert_calls.fetch_add(1, Ordering::Relaxed); + Ok(if self.create_on_call == Some(call) { + kv::StoreOutcome::Created(1) + } else { + kv::StoreOutcome::Exists(1) + }) + } + + async fn get( + &self, + _key: &kv::Key, + ) -> std::result::Result, kv::StoreError> { + Ok(None) + } + + async fn delete(&self, _key: &kv::Key) -> std::result::Result<(), kv::StoreError> { + Ok(()) + } + + async fn watch( + &self, + ) -> std::result::Result< + Pin + Send + '_>>, + kv::StoreError, + > { + Ok(Box::pin(futures::stream::pending())) + } + + async fn entries( + &self, + ) -> std::result::Result, kv::StoreError> { + Ok(HashMap::new()) + } + } + + #[tokio::test] + async fn existing_claim_does_not_poll_proposal() { + let client = KVStoreDiscovery::new(kv::Manager::memory(), CancellationToken::new()); + let mut first: ClaimPayloadFuture<'_> = Box::pin(async { Ok(payload(7)) }); + assert_eq!( + client + .create_or_get_claim("scope/session", &mut first) + .await + .unwrap(), + ClaimOutcome::Created(payload(7)) + ); + + let polled = Arc::new(AtomicBool::new(false)); + let proposal_polled = polled.clone(); + let mut second: ClaimPayloadFuture<'_> = Box::pin(async move { + proposal_polled.store(true, Ordering::Relaxed); + Ok(payload(8)) + }); + assert_eq!( + client + .create_or_get_claim("scope/session", &mut second) + .await + .unwrap(), + ClaimOutcome::Existing(payload(7)) + ); + assert!(!polled.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn competing_claims_return_one_created_winner() { + let client = Arc::new(KVStoreDiscovery::new( + kv::Manager::memory(), + CancellationToken::new(), + )); + let barrier = Arc::new(tokio::sync::Barrier::new(8)); + let mut tasks = Vec::new(); + for worker_id in 0..8 { + let client = client.clone(); + let barrier = barrier.clone(); + tasks.push(tokio::spawn(async move { + let mut proposal: ClaimPayloadFuture<'_> = + Box::pin(async move { Ok(payload(worker_id)) }); + barrier.wait().await; + client + .create_or_get_claim("scope/race", &mut proposal) + .await + .unwrap() + })); + } + + let outcomes = futures::future::join_all(tasks) + .await + .into_iter() + .map(Result::unwrap) + .collect::>(); + assert_eq!( + outcomes + .iter() + .filter(|outcome| matches!(outcome, ClaimOutcome::Created(_))) + .count(), + 1 + ); + let winner = match outcomes + .iter() + .find(|outcome| matches!(outcome, ClaimOutcome::Created(_))) + .unwrap() + { + ClaimOutcome::Created(payload) => payload, + _ => unreachable!(), + }; + assert!(outcomes.iter().all(|outcome| match outcome { + ClaimOutcome::Created(payload) | ClaimOutcome::Existing(payload) => payload == winner, + ClaimOutcome::Unsupported => false, + })); + assert_eq!(client.claims.watcher_probe.starts(), 1); + } + + #[tokio::test] + async fn claim_disappearance_retries_and_is_bounded() { + let key = kv::Key::new("scope/disappearing".to_string()); + let recovering = DisappearingBucket { + insert_calls: AtomicUsize::new(0), + create_on_call: Some(1), + }; + let mut proposal: ClaimPayloadFuture<'_> = Box::pin(async { Ok(payload(7)) }); + assert_eq!( + KVStoreDiscovery::create_or_get_in_bucket(&recovering, &key, &mut proposal) + .await + .unwrap(), + ClaimOutcome::Created(payload(7)) + ); + assert_eq!(recovering.insert_calls.load(Ordering::Relaxed), 2); + + let exhausting = DisappearingBucket { + insert_calls: AtomicUsize::new(0), + create_on_call: None, + }; + let mut proposal: ClaimPayloadFuture<'_> = Box::pin(async { Ok(payload(8)) }); + let error = KVStoreDiscovery::create_or_get_in_bucket(&exhausting, &key, &mut proposal) + .await + .unwrap_err(); + assert!(error.to_string().contains("3 competing insert attempts")); + assert_eq!(exhausting.insert_calls.load(Ordering::Relaxed), 3); + } + + #[tokio::test] + async fn claim_watcher_ignores_put_emits_delete_and_close_is_idempotent() { + let cancel = CancellationToken::new(); + let client = KVStoreDiscovery::new(kv::Manager::memory(), cancel.clone()); + let mut events = client.subscribe_claim_events().unwrap(); + let mut proposal: ClaimPayloadFuture<'_> = Box::pin(async { Ok(payload(7)) }); + client + .create_or_get_claim("scope/close", &mut proposal) + .await + .unwrap(); + + assert_eq!(client.claims.watcher_probe.starts(), 1); + assert_eq!(client.claims.watcher_probe.active(), 1); + + assert_eq!( + client.close_claim("scope/close").await.unwrap(), + ClaimCloseOutcome::Closed + ); + assert_eq!( + events.recv().await.unwrap(), + ClaimEvent::Delete("scope/close".to_string()) + ); + assert_eq!( + client.close_claim("scope/close").await.unwrap(), + ClaimCloseOutcome::Closed + ); + + cancel.cancel(); + } + + #[tokio::test] + async fn claim_watcher_stops_on_cancellation() { + let store = Arc::new(kv::Manager::memory()); + let cancel = CancellationToken::new(); + let (events, _) = broadcast::channel(16); + let (ready_tx, ready_rx) = oneshot::channel(); + let watcher = tokio::spawn(KVStoreDiscovery::run_claim_watcher( + store, + cancel.clone(), + events, + ready_tx, + )); + ready_rx.await.unwrap().unwrap(); + + cancel.cancel(); + tokio::time::timeout(Duration::from_secs(1), watcher) + .await + .expect("claim watcher did not stop after cancellation") + .unwrap(); + } + #[tokio::test] async fn test_kv_store_discovery_register_endpoint() { let store = kv::Manager::memory(); diff --git a/lib/runtime/src/discovery/mod.rs b/lib/runtime/src/discovery/mod.rs index d4a03715f4a3..91ec2b76c911 100644 --- a/lib/runtime/src/discovery/mod.rs +++ b/lib/runtime/src/discovery/mod.rs @@ -5,7 +5,9 @@ use anyhow::{Context, Result}; use async_trait::async_trait; use futures::Stream; use serde::{Deserialize, Serialize}; +use std::future::Future; use std::pin::Pin; +use tokio::sync::broadcast; use tokio_util::sync::CancellationToken; mod metadata; @@ -23,6 +25,28 @@ pub mod utils; use crate::component::{DeviceType, TransportType}; pub use utils::watch_and_extract_field; +pub type ClaimPayload = serde_json::Value; +pub type ClaimPayloadFuture<'a> = Pin> + Send + 'a>>; + +#[derive(Debug, Clone, PartialEq)] +pub enum ClaimOutcome { + Created(ClaimPayload), + Existing(ClaimPayload), + Unsupported, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ClaimCloseOutcome { + Closed, + Unsupported, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ClaimEvent { + Delete(String), + Reset, +} + /// Transport kind for event plane - used for configuration and env var selection. /// /// This enum represents the *type* of transport without connection details. @@ -844,6 +868,34 @@ pub trait Discovery: Send + Sync { cancel_token: Option, ) -> Result; + /// Returns an existing immutable claim or atomically creates it from a deferred proposal. + /// + /// Implementations read before polling `proposed_payload`, atomically insert only when + /// absent, and return the winning stored payload after an insertion race. Payloads in + /// [`ClaimOutcome::Created`] and [`ClaimOutcome::Existing`] are authoritative. + /// [`ClaimOutcome::Unsupported`] leaves coordination process-local. Storage errors must + /// propagate to the caller before scheduler bookkeeping or dispatch. + async fn create_or_get_claim( + &self, + _key: &str, + _proposed_payload: &mut ClaimPayloadFuture<'_>, + ) -> Result { + Ok(ClaimOutcome::Unsupported) + } + + /// Idempotently closes an immutable claim. + /// + /// Close is terminal under the session-ID no-reuse contract; deleting an absent claim + /// succeeds. + async fn close_claim(&self, _key: &str) -> Result { + Ok(ClaimCloseOutcome::Unsupported) + } + + /// Subscribes to process-local claim invalidation events. + fn subscribe_claim_events(&self) -> Option> { + None + } + /// Clean up resources held by this discovery backend. /// For KV store backends, this deletes owned registrations immediately rather than /// waiting for TTL expiry. Default is a no-op for backends that don't need cleanup. diff --git a/lib/runtime/src/pipeline/network/egress/push_router.rs b/lib/runtime/src/pipeline/network/egress/push_router.rs index 217739cc5c84..42dddd5aa770 100644 --- a/lib/runtime/src/pipeline/network/egress/push_router.rs +++ b/lib/runtime/src/pipeline/network/egress/push_router.rs @@ -825,6 +825,32 @@ where Ok((metadata, stream)) } + /// Book a previously arbitrated worker and dispatch without reselection or fallback. + pub async fn book_and_dispatch_exact( + &self, + mut request: SingleIn, + instance_id: u64, + advance_round_robin: bool, + prepare: F, + ) -> anyhow::Result<(M, ManyOut)> + where + F: FnOnce(&mut T, u64) -> anyhow::Result, + { + if advance_round_robin && self.router_mode == RouterMode::RoundRobin { + self.round_robin_counter.fetch_add(1, Ordering::Relaxed); + } + let (instance_id, permit) = self + .select_exact_target(request.content(), Some(instance_id)) + .await?; + let metadata = prepare(&mut request, instance_id)?; + let stream = self.dispatch_exact(request, instance_id).await?; + let stream = match permit { + Some(permit) => permit.into_tracked_stream(stream), + None => stream, + }; + Ok((metadata, stream)) + } + /// Issue a request using device-aware weighted routing. /// /// Instances are partitioned by device type (CPU vs non-CPU), then the router @@ -1075,6 +1101,23 @@ where } } + /// Peek the worker this routing mode would choose for a request without booking it. + pub fn peek_worker_for_request(&self, request: &T) -> Option { + let instance_ids = self.client.routing_instances().free_ids().to_vec(); + if instance_ids.is_empty() { + return None; + } + + match self.router_mode { + RouterMode::DeviceAwareWeighted => { + let state = self.occupancy_state.as_deref()?; + let selection = self.device_aware_candidates(request, state, &instance_ids); + state.peek_min(&selection.candidates) + } + _ => self.peek_next_worker(), + } + } + async fn select_exact_target( &self, request: &T, diff --git a/lib/runtime/src/storage/kv.rs b/lib/runtime/src/storage/kv.rs index 544a056e27f1..207142842b50 100644 --- a/lib/runtime/src/storage/kv.rs +++ b/lib/runtime/src/storage/kv.rs @@ -273,6 +273,10 @@ impl Manager { Manager(Arc::new(s)) } + pub fn is_memory(&self) -> bool { + matches!(self.0.as_ref(), KeyValueStoreEnum::Memory(_)) + } + pub async fn get_or_create_bucket( &self, bucket_name: &str, diff --git a/tests/router/common.py b/tests/router/common.py index 61c42aaba46f..e62f6d158776 100644 --- a/tests/router/common.py +++ b/tests/router/common.py @@ -9,6 +9,7 @@ import random import threading import time +import uuid from typing import TYPE_CHECKING, Any, Callable, Optional import aiohttp @@ -521,6 +522,208 @@ async def verify_consumer_lifecycle(): kv_router.__exit__(None, None, None) +def _test_distributed_session_affinity( + engine_workers, + block_size: int, + request, + router_ports: list[int], + test_payload: dict[str, Any], + store_backend: str = "etcd", +): + """Verify shared affinity claims override conflicting KV-prefix placement.""" + with ( + FrontendRouterProcess( + request, + block_size, + router_ports[0], + engine_workers.namespace, + store_backend, + router_mode="kv", + min_initial_workers=engine_workers.num_workers, + event_plane="nats", + session_affinity_ttl_secs=300, + ) as first_router, + FrontendRouterProcess( + request, + block_size, + router_ports[1], + engine_workers.namespace, + store_backend, + router_mode="kv", + min_initial_workers=engine_workers.num_workers, + event_plane="nats", + session_affinity_ttl_secs=300, + ) as second_router, + ): + urls = [f"http://localhost:{port}/v1/chat/completions" for port in router_ports] + + async def run_test() -> None: + runtime = get_runtime(store_backend, "nats") + endpoint = runtime.endpoint( + f"{engine_workers.namespace}.{engine_workers.component_name}.generate" + ) + worker_ids = sorted( + await poll_for_worker_instances(endpoint, engine_workers.num_workers) + ) + assert len(worker_ids) >= 2 + worker_a, worker_b = worker_ids[:2] + + for port in router_ports: + await wait_for_frontend_ready( + frontend_url=f"http://localhost:{port}", + expected_num_workers=engine_workers.num_workers, + timeout=120, + engine_workers=engine_workers, + store_backend=store_backend, + request_plane="nats", + ) + + suffix = uuid.uuid4().hex + prefix_a = " ".join([f"affinity-alpha-{suffix}"] * (block_size * 2)) + prefix_b = " ".join([f"affinity-beta-{suffix}"] * (block_size * 2)) + session_a = f"distributed-affinity-a-{uuid.uuid4()}" + session_b = f"distributed-affinity-b-{uuid.uuid4()}" + + def payload(content: str, *, query_only: bool = False) -> dict[str, Any]: + annotations = ["query_instance_id:"] if query_only else [] + return { + **test_payload, + "messages": [{"role": "user", "content": content}], + "stream": True, + "max_tokens": 1, + "nvext": { + "annotations": annotations, + "extra_fields": ["worker_id"], + }, + } + + async def send( + client: aiohttp.ClientSession, + url: str, + request_payload: dict[str, Any], + headers: dict[str, str] | None = None, + ) -> tuple[int, int]: + async with client.post( + url, json=request_payload, headers=headers + ) as response: + body = await response.text() + assert response.status == 200, body + + worker_info = None + for line in body.splitlines(): + if not line.startswith("data:"): + continue + data = line[5:].strip() + if data == "[DONE]": + continue + candidate = json.loads(data).get("nvext", {}).get("worker_id") + if candidate: + worker_info = candidate + + assert worker_info is not None, body + return ( + worker_info["decode_worker_id"], + worker_info["decode_dp_rank"], + ) + + async def wait_for_prefix_target( + client: aiohttp.ClientSession, + url: str, + content: str, + expected: tuple[int, int], + ) -> None: + for _ in range(50): + if ( + await send(client, url, payload(content, query_only=True)) + == expected + ): + return + await asyncio.sleep(0.1) + raise AssertionError( + f"KV events did not make prefix target {expected} visible" + ) + + session_a_headers = {"x-dynamo-session-id": session_a} + session_b_headers = {"x-dynamo-session-id": session_b} + proposal_a = { + **session_a_headers, + "x-dynamo-worker-instance-id": str(worker_a), + "x-dynamo-dp-rank": "0", + } + proposal_b = { + **session_b_headers, + "x-dynamo-worker-instance-id": str(worker_b), + "x-dynamo-dp-rank": "0", + } + + async with aiohttp.ClientSession() as client: + assert await send(client, urls[0], payload(prefix_a), proposal_a) == ( + worker_a, + 0, + ) + assert await send(client, urls[1], payload(prefix_b), proposal_b) == ( + worker_b, + 0, + ) + + await wait_for_prefix_target(client, urls[0], prefix_a, (worker_a, 0)) + await wait_for_prefix_target(client, urls[1], prefix_b, (worker_b, 0)) + + assert await send( + client, urls[0], payload(prefix_a), session_b_headers + ) == (worker_b, 0) + assert await send( + client, urls[1], payload(prefix_b), session_a_headers + ) == (worker_a, 0) + + first_evictions = first_router.read_logs().count( + "evicted session affinity cache entry" + ) + second_evictions = second_router.read_logs().count( + "evicted session affinity cache entry" + ) + assert await send( + client, + urls[1], + payload(prefix_b), + { + **session_a_headers, + "x-dynamo-session-final": "true", + }, + ) == (worker_a, 0) + assert await send( + client, + urls[0], + payload(prefix_a), + { + **session_b_headers, + "x-dynamo-session-final": "true", + }, + ) == (worker_b, 0) + + for _ in range(50): + first_observed = ( + first_router.read_logs().count( + "evicted session affinity cache entry" + ) + > first_evictions + ) + second_observed = ( + second_router.read_logs().count( + "evicted session affinity cache entry" + ) + > second_evictions + ) + if first_observed and second_observed: + return + await asyncio.sleep(0.1) + raise AssertionError( + "frontends did not observe terminal claim deletion invalidation" + ) + + asyncio.run(run_test()) + + def _test_remote_indexer_decisions( engine_workers, model_name: str, diff --git a/tests/router/router_process.py b/tests/router/router_process.py index 330453eb5081..988fbbf19886 100644 --- a/tests/router/router_process.py +++ b/tests/router/router_process.py @@ -58,6 +58,7 @@ def __init__( serve_indexer: bool = False, use_remote_indexer: bool = False, event_plane: str | None = None, + session_affinity_ttl_secs: int | None = None, ): command = [ sys.executable, @@ -102,6 +103,11 @@ def __init__( if use_remote_indexer: command.append("--use-remote-indexer") + if session_affinity_ttl_secs is not None: + command.extend( + ["--router-session-affinity-ttl-secs", str(session_affinity_ttl_secs)] + ) + if router_aic_config is not None: command.extend( [ @@ -146,7 +152,7 @@ def __init__( ], log_dir=request.node.name, terminate_all_matching_process_names=False, - display_name=f"dynamo-frontend-{router_mode}", + display_name=f"dynamo-frontend-{router_mode}-{frontend_port}", ) self.port = frontend_port self.router_mode = router_mode diff --git a/tests/router/test_router_e2e_with_mockers.py b/tests/router/test_router_e2e_with_mockers.py index 273b58788bd4..c848bade8a5f 100644 --- a/tests/router/test_router_e2e_with_mockers.py +++ b/tests/router/test_router_e2e_with_mockers.py @@ -23,6 +23,7 @@ _test_disagg_direct_mode, _test_disagg_router_overload_529, _test_disagg_topology_required_prefill_pin_match_and_mismatch, + _test_distributed_session_affinity, _test_python_router_bindings, _test_remote_indexer_decisions, _test_router_decisions_disagg_round_robin_prefill_dp_rank, @@ -490,6 +491,44 @@ def test_mocker_two_kv_router( ) +@pytest.mark.parametrize("store_backend", ["etcd", "file"]) +@pytest.mark.timeout(180) +def test_mocker_distributed_session_affinity( + request, + runtime_services_dynamic_ports, + predownload_tokenizers, + file_storage_backend, + store_backend, + monkeypatch, +): + """Shared claims override conflicting KV-prefix routing on another frontend.""" + current_log = os.environ.get("DYN_LOG", "info") + monkeypatch.setenv( + "DYN_LOG", + f"{current_log},dynamo_llm::session_affinity::coordinator=debug", + ) + mocker_args = { + "speedup_ratio": SPEEDUP_RATIO, + "block_size": BLOCK_SIZE, + "durable_kv_events": False, + } + + with MockerProcess( + request, + mocker_args=mocker_args, + num_mockers=NUM_MOCKERS, + store_backend=store_backend, + ) as mockers: + _test_distributed_session_affinity( + engine_workers=mockers, + block_size=BLOCK_SIZE, + request=request, + router_ports=allocate_frontend_ports(request, 2), + test_payload=TEST_PAYLOAD, + store_backend=store_backend, + ) + + @pytest.mark.parametrize( "durable_kv_events", [False], ids=["nondurable"], indirect=True ) # Use NATS Core (local indexer) From 065abee5f1214c8a4b408975d872eab0a6c4f8e6 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Mon, 29 Jun 2026 22:32:37 -0700 Subject: [PATCH 2/6] fix(router): harden affinity invalidation and arbitration Signed-off-by: PeaBrane --- Cargo.lock | 10 -- lib/bindings/python/Cargo.lock | 10 -- lib/llm/src/kv_router/push_router.rs | 43 +++++---- lib/llm/src/session_affinity/coordinator.rs | 24 +++-- lib/llm/src/session_affinity/push_router.rs | 101 +++++++++++--------- lib/llm/src/session_affinity/tests.rs | 74 +++++++++----- lib/runtime/Cargo.toml | 3 +- lib/runtime/src/discovery/kv_store.rs | 78 ++++++++++++--- lib/runtime/src/storage/kv/file.rs | 92 ++++++++++++++++-- tests/router/common.py | 34 ++++--- 10 files changed, 317 insertions(+), 152 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ddba6950726e..7c2eb88d945c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3433,15 +3433,6 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" -[[package]] -name = "fsevent-sys" -version = "4.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ee7a02da4d231650c7cea31349b889be2f45ddb3ef3032d2ec8185f6313fd2" -dependencies = [ - "libc", -] - [[package]] name = "futures" version = "0.3.32" @@ -5776,7 +5767,6 @@ checksum = "6205bd8bb1e454ad2e27422015fb5e4f2bcc7e08fa8f27058670d208324a4d2d" dependencies = [ "bitflags 2.11.1", "filetime", - "fsevent-sys", "inotify", "kqueue", "libc", diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index f1cbbb3c1801..710d5fc65fe8 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2994,15 +2994,6 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" -[[package]] -name = "fsevent-sys" -version = "4.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ee7a02da4d231650c7cea31349b889be2f45ddb3ef3032d2ec8185f6313fd2" -dependencies = [ - "libc", -] - [[package]] name = "futures" version = "0.3.32" @@ -5096,7 +5087,6 @@ checksum = "6205bd8bb1e454ad2e27422015fb5e4f2bcc7e08fa8f27058670d208324a4d2d" dependencies = [ "bitflags 2.13.0", "filetime", - "fsevent-sys", "inotify", "kqueue", "libc", diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index 33f43326127a..3fa261cf94a1 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -156,15 +156,18 @@ impl KvPushRouter { let operation = affinity .acquire_with_context(&session_id, request_context.as_ref()) .await?; - let proposed_payload: ClaimPayloadFuture<'_> = Box::pin(async { - let selection = self.select_request(request, phase, true, None).await?; - let target = AffinityTarget { - worker_id: selection.instance_id, - dp_rank: Some(selection.dp_rank), - }; - Ok(serde_json::to_value(target)?) - }); - let resolved = operation.resolve(proposed_payload).await?; + let resolved = operation + .resolve(|| -> ClaimPayloadFuture<'_> { + Box::pin(async { + let selection = self.select_request(request, phase, true, None).await?; + let target = AffinityTarget { + worker_id: selection.instance_id, + dp_rank: Some(selection.dp_rank), + }; + Ok(serde_json::to_value(target)?) + }) + }) + .await?; let worker = affinity_worker(resolved.target()); let selection = self.select_request(request, phase, false, worker).await?; Ok((selection, Some(resolved))) @@ -663,9 +666,7 @@ mod tests { .acquire(&session_id) .await .unwrap() - .resolve(Box::pin(async move { - Ok(serde_json::to_value(original_target)?) - })) + .resolve(|| Box::pin(async move { Ok(serde_json::to_value(original_target)?) })) .await .unwrap(); drop(resolved); @@ -703,12 +704,14 @@ mod tests { .acquire(&session_id) .await .unwrap() - .resolve(Box::pin(async move { - Ok(serde_json::to_value(AffinityTarget { - worker_id: 8, - dp_rank: Some(0), - })?) - })) + .resolve(|| { + Box::pin(async move { + Ok(serde_json::to_value(AffinityTarget { + worker_id: 8, + dp_rank: Some(0), + })?) + }) + }) .await .unwrap(); assert_eq!(resolved.target(), original_target); @@ -732,9 +735,7 @@ mod tests { .acquire(&session_id) .await .unwrap() - .resolve(Box::pin( - async move { Ok(serde_json::to_value(bound_target)?) }, - )) + .resolve(|| Box::pin(async move { Ok(serde_json::to_value(bound_target)?) })) .await .unwrap(); drop(resolved); diff --git a/lib/llm/src/session_affinity/coordinator.rs b/lib/llm/src/session_affinity/coordinator.rs index 8dea2e88f599..f981fef307da 100644 --- a/lib/llm/src/session_affinity/coordinator.rs +++ b/lib/llm/src/session_affinity/coordinator.rs @@ -259,13 +259,15 @@ impl AffinityCoordinator { } fn clear_entries(inner: &AffinityCoordinatorInner) { - for entry in inner.entries.iter() { - if let AffinityEntry::Initializing { notify, .. } = entry.value() { + let mut removed = 0; + inner.entries.retain(|_, entry| { + if let AffinityEntry::Initializing { notify, .. } = entry { notify.notify_waiters(); } - } - inner.entries.clear(); - inner.entry_count.store(0, Ordering::Relaxed); + removed += 1; + false + }); + Self::decrement_entry_count(inner, removed); tracing::debug!("cleared session affinity cache after claim watcher reset"); } @@ -552,12 +554,12 @@ pub(crate) enum AffinityAcquire { } impl AffinityAcquire { - pub(crate) async fn resolve( - self, - proposed_payload: ClaimPayloadFuture<'_>, - ) -> Result { + pub(crate) async fn resolve<'a, F>(self, proposed_payload: F) -> Result + where + F: FnOnce() -> ClaimPayloadFuture<'a> + Send, + { match self { - Self::Initialize(initialization) => initialization.resolve(proposed_payload).await, + Self::Initialize(initialization) => initialization.resolve(proposed_payload()).await, Self::Bound { target, lease } => Ok(ResolvedAffinity { target, lease, @@ -750,6 +752,8 @@ impl CloseAction { return; }; let claim_key = self.claim_key; + // TODO: Drive backend close to completion before returning stream EOF. This detached + // task keeps early stream drops best-effort and can be cancelled during runtime shutdown. tokio::spawn(async move { if let Err(error) = claims.close(&claim_key).await { tracing::error!(%claim_key, %error, "failed to close session affinity claim"); diff --git a/lib/llm/src/session_affinity/push_router.rs b/lib/llm/src/session_affinity/push_router.rs index 4d144187691d..fe68be1b476b 100644 --- a/lib/llm/src/session_affinity/push_router.rs +++ b/lib/llm/src/session_affinity/push_router.rs @@ -77,10 +77,10 @@ impl SessionAffinityPushRouter { async fn resolve_affinity( &self, session_id: &crate::protocols::common::extensions::SessionAffinityId, - explicit: Option, + phase: RequestPhase, request: &PreprocessedRequest, request_context: &dyn AsyncEngineContext, - ) -> Result { + ) -> Result<(ResolvedAffinity, bool), Error> { let affinity = self .affinity .as_ref() @@ -88,26 +88,37 @@ impl SessionAffinityPushRouter { let operation = affinity .acquire_with_context(session_id, request_context) .await?; - let proposed_payload: ClaimPayloadFuture<'_> = Box::pin(async move { - let target = explicit - .or_else(|| { - self.inner - .peek_worker_for_request(request) - .map(|worker_id| AffinityTarget { - worker_id, - dp_rank: None, + let resolved = operation + .resolve(|| -> ClaimPayloadFuture<'_> { + Box::pin(async move { + let target = explicit_target(request, phase)? + .or_else(|| { + self.inner + .peek_worker_for_request(request) + .map(|worker_id| AffinityTarget { + worker_id, + dp_rank: None, + }) }) + .ok_or_else(|| { + if self.direct { + invalid_argument( + "worker ID required to create Direct session affinity", + ) + } else { + anyhow::anyhow!("no worker available for session affinity") + } + })?; + Ok(serde_json::to_value(target)?) }) - .ok_or_else(|| { - if self.direct { - invalid_argument("worker ID required to create Direct session affinity") - } else { - anyhow::anyhow!("no worker available for session affinity") - } - })?; - Ok(serde_json::to_value(target)?) - }); - operation.resolve(proposed_payload).await + }) + .await?; + let proposal_was_explicit = if resolved.was_created() { + explicit_target(request, phase)?.is_some() + } else { + false + }; + Ok((resolved, proposal_was_explicit)) } pub async fn select_and_dispatch_prefill( @@ -132,8 +143,8 @@ impl SessionAffinityPushRouter { }) .await; } - let explicit = explicit_target(&request, RequestPhase::Prefill)?; let Some(session_id) = session_id else { + let explicit = explicit_target(&request, RequestPhase::Prefill)?; let Some(pinned_worker) = explicit else { return Err(invalid_argument( "worker ID required for prefill request in Direct routing mode", @@ -150,12 +161,15 @@ impl SessionAffinityPushRouter { }; let is_query_only = request.get_annotation_value("query_instance_id").is_some(); if is_query_only { - let selected = self + let bound = self .affinity .as_ref() .expect("affinity query requires an enabled coordinator") - .query_target(&session_id)? - .or(explicit); + .query_target(&session_id)?; + let selected = match bound { + Some(target) => Some(target), + None => explicit_target(&request, RequestPhase::Prefill)?, + }; let rank = selected.and_then(|target| target.dp_rank); return self .inner @@ -176,17 +190,17 @@ impl SessionAffinityPushRouter { let close_on_finish = session_final(request.content()); let request_context = request.context(); - let resolved = self + let (resolved, proposal_was_explicit) = self .resolve_affinity( &session_id, - explicit, + RequestPhase::Prefill, request.content(), request_context.as_ref(), ) .await?; let target = resolved.target(); let rank = target.dp_rank; - let advance_round_robin = resolved.was_created() && explicit.is_none(); + let advance_round_robin = resolved.was_created() && !proposal_was_explicit; let ((metadata, _), stream) = self .inner .book_and_dispatch_exact( @@ -224,8 +238,8 @@ impl AsyncEngine, ManyOut, Error> if !self.direct && session_id.is_none() { return self.inner.generate(request).await; } - let explicit = explicit_target(&request, phase)?; let Some(session_id) = session_id else { + let explicit = explicit_target(&request, phase)?; let Some(target) = explicit else { return Err(invalid_argument(format!( "worker ID required for {phase} request in Direct routing mode" @@ -236,12 +250,15 @@ impl AsyncEngine, ManyOut, Error> let is_query_only = request.get_annotation_value("query_instance_id").is_some(); if is_query_only { - let target = self + let bound = self .affinity .as_ref() .expect("affinity query requires an enabled coordinator") - .query_target(&session_id)? - .or(explicit); + .query_target(&session_id)?; + let target = match bound { + Some(target) => Some(target), + None => explicit_target(&request, phase)?, + }; let rank = target.and_then(|target| target.dp_rank); let (_, stream) = self .inner @@ -268,17 +285,17 @@ impl AsyncEngine, ManyOut, Error> let close_on_finish = session_final(request.content()); let request_context = request.context(); - let resolved = self + let (resolved, proposal_was_explicit) = self .resolve_affinity( &session_id, - explicit, + phase, request.content(), request_context.as_ref(), ) .await?; let target = resolved.target(); let rank = target.dp_rank; - let advance_round_robin = resolved.was_created() && explicit.is_none(); + let advance_round_robin = resolved.was_created() && !proposal_was_explicit; let (_, stream) = self .inner .book_and_dispatch_exact( @@ -501,7 +518,7 @@ mod tests { } #[tokio::test] - async fn session_affinity_unavailable_target_remains_bound() { + async fn session_affinity_binding_wins_before_invalid_explicit_proposal() { let runtime = Runtime::from_current().unwrap(); let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) @@ -529,19 +546,15 @@ mod tests { .acquire(&session_id) .await .unwrap() - .resolve(Box::pin(async move { - Ok(serde_json::to_value(unavailable_target)?) - })) + .resolve(|| Box::pin(async move { Ok(serde_json::to_value(unavailable_target)?) })) .await .unwrap(); drop(resolved); - assert!( - router - .generate(affinity_request(None, false)) - .await - .is_err() - ); + let mut request = affinity_request(None, false); + request.routing_mut().dp_rank = Some(0); + let error = router.generate(request).await.unwrap_err(); + assert!(!error.to_string().contains("DP rank requires")); assert_eq!( affinity(&router).query_target(&session_id).unwrap(), Some(AffinityTarget { diff --git a/lib/llm/src/session_affinity/tests.rs b/lib/llm/src/session_affinity/tests.rs index 8bdd1212f240..7957109c1c8c 100644 --- a/lib/llm/src/session_affinity/tests.rs +++ b/lib/llm/src/session_affinity/tests.rs @@ -163,7 +163,7 @@ async fn resolve_local( .acquire(session_id) .await .unwrap() - .resolve(target_payload(selected)) + .resolve(|| target_payload(selected)) .await .unwrap() } @@ -290,14 +290,14 @@ async fn session_affinity_initialization_is_atomic() { assert!(!waiter.is_finished()); let first = first - .resolve(target_payload(target(7, Some(0)))) + .resolve(|| target_payload(target(7, Some(0)))) .await .unwrap(); let second = waiter .await .unwrap() .unwrap() - .resolve(target_payload(target(8, Some(0)))) + .resolve(|| target_payload(target(8, Some(0)))) .await .unwrap(); assert_eq!(second.target(), target(7, Some(0))); @@ -316,7 +316,7 @@ async fn distributed_existing_winner_skips_proposal_and_cache_hits_skip_discover .acquire(&session_id) .await .unwrap() - .resolve(target_payload(target(7, Some(0)))) + .resolve(|| target_payload(target(7, Some(0)))) .await .unwrap(); assert_eq!(created.target(), target(7, Some(0))); @@ -325,15 +325,16 @@ async fn distributed_existing_winner_skips_proposal_and_cache_hits_skip_discover let proposal_polled = Arc::new(AtomicBool::new(false)); let polled = proposal_polled.clone(); - let proposed: ClaimPayloadFuture<'_> = Box::pin(async move { - polled.store(true, Ordering::Relaxed); - Ok(serde_json::to_value(target(8, Some(0)))?) - }); let winner = second .acquire(&session_id) .await .unwrap() - .resolve(proposed) + .resolve(|| { + Box::pin(async move { + polled.store(true, Ordering::Relaxed); + Ok(serde_json::to_value(target(8, Some(0)))?) + }) + }) .await .unwrap(); assert_eq!(winner.target(), target(7, Some(0))); @@ -342,21 +343,20 @@ async fn distributed_existing_winner_skips_proposal_and_cache_hits_skip_discover drop(winner); assert_eq!(discovery.create_calls.load(Ordering::Relaxed), 2); - let cached_proposal_polled = Arc::new(AtomicBool::new(false)); - let polled = cached_proposal_polled.clone(); - let proposed: ClaimPayloadFuture<'_> = Box::pin(async move { - polled.store(true, Ordering::Relaxed); - Ok(serde_json::to_value(target(9, Some(0)))?) - }); + let cached_proposal_constructed = Arc::new(AtomicBool::new(false)); + let constructed = cached_proposal_constructed.clone(); let cached = second .acquire(&session_id) .await .unwrap() - .resolve(proposed) + .resolve(|| { + constructed.store(true, Ordering::Relaxed); + target_payload(target(9, Some(0))) + }) .await .unwrap(); assert_eq!(cached.target(), target(7, Some(0))); - assert!(!cached_proposal_polled.load(Ordering::Relaxed)); + assert!(!cached_proposal_constructed.load(Ordering::Relaxed)); assert_eq!(discovery.create_calls.load(Ordering::Relaxed), 2); } @@ -376,7 +376,7 @@ async fn distributed_delete_evicts_one_entry_and_duplicate_is_harmless() { .acquire(session) .await .unwrap() - .resolve(target_payload(target(worker_id, Some(0)))) + .resolve(|| target_payload(target(worker_id, Some(0)))) .await .unwrap(), ); @@ -403,7 +403,7 @@ async fn distributed_subscriber_lag_clears_all_entries() { .acquire(&session_id) .await .unwrap() - .resolve(target_payload(target(8, Some(0)))) + .resolve(|| target_payload(target(8, Some(0)))) .await .unwrap(), ); @@ -425,7 +425,7 @@ async fn distributed_reset_and_disconnect_clear_entries() { .acquire(&session_id) .await .unwrap() - .resolve(target_payload(target(8, Some(0)))) + .resolve(|| target_payload(target(8, Some(0)))) .await .unwrap(), ); @@ -437,7 +437,7 @@ async fn distributed_reset_and_disconnect_clear_entries() { .acquire(&session_id) .await .unwrap() - .resolve(target_payload(target(99, Some(0)))) + .resolve(|| target_payload(target(99, Some(0)))) .await .unwrap(), ); @@ -454,7 +454,7 @@ async fn terminal_close_evicts_synchronously() { .acquire(&session_id) .await .unwrap() - .resolve(target_payload(target(7, Some(0)))) + .resolve(|| target_payload(target(7, Some(0)))) .await .unwrap(); let key = coordinator.claim_key_for_test(&session_id); @@ -495,6 +495,34 @@ async fn session_affinity_initializer_cancellation_wakes_waiter() { )); } +#[tokio::test] +async fn distributed_reset_wakes_initializing_waiter_and_preserves_entry_count() { + let discovery = ClaimTestDiscovery::new(16); + let coordinator = distributed_coordinator(discovery.clone()); + let first = coordinator.acquire(&session_id()).await.unwrap(); + let AffinityAcquire::Initialize(first) = first else { + panic!("first request must initialize"); + }; + + let waiter_coordinator = coordinator.clone(); + let waiter = tokio::spawn(async move { waiter_coordinator.acquire(&session_id()).await }); + coordinator.wait_for_initializing_waiter().await; + discovery.emit(ClaimEvent::Reset); + + let next = tokio::time::timeout(Duration::from_secs(1), waiter) + .await + .expect("reset did not wake initializing waiter") + .unwrap() + .unwrap(); + assert!(matches!(&next, AffinityAcquire::Initialize(_))); + assert_eq!(coordinator.entry_count(), 1); + + drop(first); + assert_eq!(coordinator.entry_count(), 1); + drop(next); + assert_eq!(coordinator.entry_count(), 0); +} + #[tokio::test(start_paused = true)] async fn session_affinity_wait_stops_when_request_is_cancelled() { let coordinator = coordinator(); @@ -535,7 +563,7 @@ async fn session_affinity_existing_binding_overrides_explicit_proposals() { .acquire(&session_id()) .await .unwrap() - .resolve(target_payload(proposal)) + .resolve(|| target_payload(proposal)) .await .unwrap(); assert_eq!(resolved.target(), target(7, None)); diff --git a/lib/runtime/Cargo.toml b/lib/runtime/Cargo.toml index 7aa77c2b4151..54ca9be9e20a 100644 --- a/lib/runtime/Cargo.toml +++ b/lib/runtime/Cargo.toml @@ -78,7 +78,8 @@ bincode = { version = "1" } console-subscriber = { version = "0.4", optional = true } educe = { version = "0.6.0" } figment = { version = "0.10.19", features = ["env", "json", "toml", "test"] } -notify = { version = "6.1", default-features = false, features = ["macos_fsevent"] } +# FSEvents omits deletes for FileStore's revision-zero hard-link publication path. +notify = { version = "6.1", default-features = false, features = ["macos_kqueue"] } libc = { version = "0.2" } local-ip-address = { version = "0.6.3" } # `release_max_level_debug` compiles out `log::trace!` in release builds (a diff --git a/lib/runtime/src/discovery/kv_store.rs b/lib/runtime/src/discovery/kv_store.rs index 26f0fbec60e2..528b20ff4c9f 100644 --- a/lib/runtime/src/discovery/kv_store.rs +++ b/lib/runtime/src/discovery/kv_store.rs @@ -39,6 +39,9 @@ pub struct KVStoreDiscovery { /// instance. `Put` events never populate caches. `Delete(key)` evicts one entry, while /// watcher loss, restart, or subscriber lag produces `Reset` so coordinators clear all /// entries rather than retain potentially stale bindings. +/// +/// TODO: `Bucket::watch` cannot yet surface etcd reconnect/compaction or FileStore +/// overflow errors, so those hidden backend failures cannot be converted into `Reset`. struct ClaimState { events: broadcast::Sender, watcher_started: OnceCell<()>, @@ -905,6 +908,50 @@ mod tests { create_on_call: Option, } + struct InsertBarrierBucket { + inner: Box, + barrier: tokio::sync::Barrier, + } + + #[async_trait] + impl kv::Bucket for InsertBarrierBucket { + async fn insert( + &self, + key: &kv::Key, + value: bytes::Bytes, + revision: u64, + ) -> std::result::Result { + self.barrier.wait().await; + self.inner.insert(key, value, revision).await + } + + async fn get( + &self, + key: &kv::Key, + ) -> std::result::Result, kv::StoreError> { + self.inner.get(key).await + } + + async fn delete(&self, key: &kv::Key) -> std::result::Result<(), kv::StoreError> { + self.inner.delete(key).await + } + + async fn watch( + &self, + ) -> std::result::Result< + Pin + Send + '_>>, + kv::StoreError, + > { + self.inner.watch().await + } + + async fn entries( + &self, + ) -> std::result::Result, kv::StoreError> { + self.inner.entries().await + } + } + #[async_trait] impl kv::Bucket for DisappearingBucket { async fn insert( @@ -978,23 +1025,29 @@ mod tests { #[tokio::test] async fn competing_claims_return_one_created_winner() { - let client = Arc::new(KVStoreDiscovery::new( - kv::Manager::memory(), - CancellationToken::new(), - )); - let barrier = Arc::new(tokio::sync::Barrier::new(8)); + let store = kv::Manager::memory(); + let bucket = Arc::new(InsertBarrierBucket { + inner: store + .get_or_create_bucket(CLAIMS_BUCKET, None) + .await + .unwrap(), + barrier: tokio::sync::Barrier::new(8), + }); + let key = Arc::new(kv::Key::new("scope/race".to_string())); let mut tasks = Vec::new(); for worker_id in 0..8 { - let client = client.clone(); - let barrier = barrier.clone(); + let bucket = bucket.clone(); + let key = key.clone(); tasks.push(tokio::spawn(async move { let mut proposal: ClaimPayloadFuture<'_> = Box::pin(async move { Ok(payload(worker_id)) }); - barrier.wait().await; - client - .create_or_get_claim("scope/race", &mut proposal) - .await - .unwrap() + KVStoreDiscovery::create_or_get_in_bucket( + bucket.as_ref(), + key.as_ref(), + &mut proposal, + ) + .await + .unwrap() })); } @@ -1022,7 +1075,6 @@ mod tests { ClaimOutcome::Created(payload) | ClaimOutcome::Existing(payload) => payload == winner, ClaimOutcome::Unsupported => false, })); - assert_eq!(client.claims.watcher_probe.starts(), 1); } #[tokio::test] diff --git a/lib/runtime/src/storage/kv/file.rs b/lib/runtime/src/storage/kv/file.rs index bbb138b4cb8a..f4520ad48b87 100644 --- a/lib/runtime/src/storage/kv/file.rs +++ b/lib/runtime/src/storage/kv/file.rs @@ -195,8 +195,9 @@ pub struct Directory { impl Directory { fn new(root: PathBuf, p: PathBuf, ttl: Duration) -> Self { - // Canonicalize root to handle symlinks (e.g., /var -> /private/var on macOS) + // Keep watched paths and event paths in the same form across symlinked roots. let canonical_root = root.canonicalize().unwrap_or_else(|_| root.clone()); + let canonical_path = p.canonicalize().unwrap_or_else(|_| p.clone()); if ttl < MIN_KEEP_ALIVE { let h_ttl = humantime::format_duration(ttl); tracing::warn!(path = %p.display(), ttl = %h_ttl, "ttl is too short, increasing to {}", humantime::format_duration(MIN_KEEP_ALIVE)); @@ -204,7 +205,7 @@ impl Directory { let ttl = cmp::max(ttl, MIN_KEEP_ALIVE); Directory { root: canonical_root, - p, + p: canonical_path, ttl, owned_files: Arc::new(Mutex::new(HashSet::new())), } @@ -441,7 +442,6 @@ impl Bucket for Directory { continue; } }; - for item_path in event.paths { // Skip if the event is for the directory itself if item_path == dir { @@ -449,9 +449,7 @@ impl Bucket for Directory { continue; } - // Canonicalize paths to handle symlinks (e.g., /var -> /private/var on macOS) - // The unwrap_or_else path is for Remove case. - let canonical_item_path = item_path.canonicalize().unwrap_or_else(|_| item_path.clone()); + let canonical_item_path = canonicalize_event_path(&item_path); let key = match canonical_item_path.strip_prefix(&root) { Ok(stripped) => Key::from_url_safe(&stripped.display().to_string()), @@ -490,7 +488,7 @@ impl Bucket for Directory { let item = KeyValue::new(key, data); yield WatchEvent::Put(item); } - EventKind::Remove(event::RemoveKind::File) => { + EventKind::Remove(_) => { yield WatchEvent::Delete(key); } _ => { @@ -589,6 +587,19 @@ fn write_temp_file_at(temp_path: &Path, value: &[u8]) -> Result PathBuf { + if let Ok(canonical_path) = path.canonicalize() { + return canonical_path; + } + let (Some(parent), Some(file_name)) = (path.parent(), path.file_name()) else { + return path.to_path_buf(); + }; + let Ok(canonical_parent) = parent.canonicalize() else { + return path.to_path_buf(); + }; + canonical_parent.join(file_name) +} + // For anyhow preserve the context fn a_to_fs_err(err: anyhow::Error) -> StoreError { StoreError::FilesystemError(format!("{err:#}")) @@ -602,11 +613,78 @@ fn to_fs_err(err: E) -> StoreError { mod tests { use std::collections::HashSet; use std::fs; + use std::os::unix::fs::symlink; + use std::time::Duration; + use futures::StreamExt; use tokio_util::sync::CancellationToken; use crate::storage::kv::{Bucket as _, FileStore, Key, Store as _, StoreOutcome}; + #[test] + fn deleted_event_path_canonicalizes_existing_parent() { + let t = tempfile::tempdir().unwrap(); + let canonical_root = t.path().join("canonical"); + let bucket = canonical_root.join("v1/claims"); + fs::create_dir_all(&bucket).unwrap(); + let linked_root = t.path().join("linked"); + symlink(&canonical_root, &linked_root).unwrap(); + + assert_eq!( + super::canonicalize_event_path(&linked_root.join("v1/claims/deleted")), + canonical_root + .canonicalize() + .unwrap() + .join("v1/claims/deleted") + ); + } + + #[tokio::test] + async fn external_delete_is_observed_under_noncanonical_root() { + let t = tempfile::tempdir().unwrap(); + let watcher_cancel = CancellationToken::new(); + let creator_cancel = CancellationToken::new(); + let watcher_store = FileStore::new(watcher_cancel.clone(), t.path()); + let creator_store = FileStore::new(creator_cancel.clone(), t.path()); + let watcher_bucket = watcher_store + .get_or_create_bucket("v1/claims", None) + .await + .unwrap(); + let creator_bucket = creator_store + .get_or_create_bucket("v1/claims", None) + .await + .unwrap(); + let mut events = watcher_bucket.watch().await.unwrap(); + let key = Key::new("scope/session".to_string()); + + creator_bucket + .insert(&key, "value".into(), 0) + .await + .unwrap(); + loop { + let event = tokio::time::timeout(Duration::from_secs(2), events.next()) + .await + .expect("FileStore watcher did not observe claim creation") + .expect("FileStore watcher ended after claim creation"); + if matches!(event, super::WatchEvent::Put(ref item) if item.key_str() == "v1/claims/scope/session") + { + break; + } + } + + creator_bucket.delete(&key).await.unwrap(); + let event = tokio::time::timeout(Duration::from_secs(2), events.next()) + .await + .expect("FileStore watcher did not observe claim deletion") + .expect("FileStore watcher ended after claim deletion"); + assert!( + matches!(event, super::WatchEvent::Delete(ref deleted) if deleted == &Key::new("v1/claims/scope/session".to_string())) + ); + + watcher_cancel.cancel(); + creator_cancel.cancel(); + } + #[tokio::test] async fn test_entries_full_path() { let t = tempfile::tempdir().unwrap(); diff --git a/tests/router/common.py b/tests/router/common.py index e62f6d158776..2222295b676c 100644 --- a/tests/router/common.py +++ b/tests/router/common.py @@ -679,9 +679,6 @@ async def wait_for_prefix_target( first_evictions = first_router.read_logs().count( "evicted session affinity cache entry" ) - second_evictions = second_router.read_logs().count( - "evicted session affinity cache entry" - ) assert await send( client, urls[1], @@ -691,6 +688,24 @@ async def wait_for_prefix_target( "x-dynamo-session-final": "true", }, ) == (worker_a, 0) + + for _ in range(50): + if ( + first_router.read_logs().count( + "evicted session affinity cache entry" + ) + > first_evictions + ): + break + await asyncio.sleep(0.1) + else: + raise AssertionError( + "first frontend did not observe session A claim deletion" + ) + + second_evictions = second_router.read_logs().count( + "evicted session affinity cache entry" + ) assert await send( client, urls[0], @@ -702,23 +717,16 @@ async def wait_for_prefix_target( ) == (worker_b, 0) for _ in range(50): - first_observed = ( - first_router.read_logs().count( - "evicted session affinity cache entry" - ) - > first_evictions - ) - second_observed = ( + if ( second_router.read_logs().count( "evicted session affinity cache entry" ) > second_evictions - ) - if first_observed and second_observed: + ): return await asyncio.sleep(0.1) raise AssertionError( - "frontends did not observe terminal claim deletion invalidation" + "second frontend did not observe session B claim deletion" ) asyncio.run(run_test()) From 4e6b45de2a3a5344541519a2fa07d11163d44c83 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Mon, 29 Jun 2026 22:42:10 -0700 Subject: [PATCH 3/6] fix(ci): refresh kvbm lockfile for notify Signed-off-by: PeaBrane --- lib/bindings/kvbm/Cargo.lock | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/lib/bindings/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index bf75e62b27d6..3b84de82f913 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -2252,15 +2252,6 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" -[[package]] -name = "fsevent-sys" -version = "4.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ee7a02da4d231650c7cea31349b889be2f45ddb3ef3032d2ec8185f6313fd2" -dependencies = [ - "libc", -] - [[package]] name = "futures" version = "0.3.32" @@ -4066,7 +4057,6 @@ checksum = "6205bd8bb1e454ad2e27422015fb5e4f2bcc7e08fa8f27058670d208324a4d2d" dependencies = [ "bitflags 2.12.1", "filetime", - "fsevent-sys", "inotify", "kqueue", "libc", From a11ec57fb1e94f03d4afdbd3b0b8c72a6598ea18 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Mon, 29 Jun 2026 22:50:11 -0700 Subject: [PATCH 4/6] fix(router): address affinity review feedback Signed-off-by: PeaBrane --- lib/llm/src/session_affinity/push_router.rs | 56 +++++++++++++++++++-- lib/llm/src/session_affinity/tests.rs | 19 +++---- lib/runtime/src/storage/kv/file.rs | 8 ++- 3 files changed, 69 insertions(+), 14 deletions(-) diff --git a/lib/llm/src/session_affinity/push_router.rs b/lib/llm/src/session_affinity/push_router.rs index fe68be1b476b..429bfb25a81c 100644 --- a/lib/llm/src/session_affinity/push_router.rs +++ b/lib/llm/src/session_affinity/push_router.rs @@ -145,17 +145,18 @@ impl SessionAffinityPushRouter { } let Some(session_id) = session_id else { let explicit = explicit_target(&request, RequestPhase::Prefill)?; - let Some(pinned_worker) = explicit else { + let Some(target) = explicit else { return Err(invalid_argument( "worker ID required for prefill request in Direct routing mode", )); }; + let rank = target.dp_rank; return self .inner .select_and_dispatch_exact( request, - Some(pinned_worker.worker_id), - move |request, worker_id| prepare(request, worker_id, None), + Some(target.worker_id), + move |request, worker_id| prepare(request, worker_id, rank), ) .await; }; @@ -330,6 +331,8 @@ fn phase_worker_id(request: &PreprocessedRequest, phase: RequestPhase) -> Option #[cfg(test)] mod tests { + use std::sync::{Arc, Mutex}; + use dynamo_runtime::{ DistributedRuntime, Runtime, distributed::DistributedConfig, @@ -517,6 +520,53 @@ mod tests { runtime.shutdown(); } + #[tokio::test] + async fn direct_prefill_without_session_preserves_explicit_rank_zero() { + let runtime = Runtime::from_current().unwrap(); + let distributed = + DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) + .await + .unwrap(); + let endpoint = distributed + .namespace("session_affinity_direct_prefill_rank".to_string()) + .unwrap() + .component("workers".to_string()) + .unwrap() + .endpoint("prefill".to_string()); + let client = endpoint.client().await.unwrap(); + endpoint.register_endpoint_instance().await.unwrap(); + let worker_id = client.wait_for_instances().await.unwrap()[0].id(); + let inner = PushRouter::from_client(client, RouterMode::Direct) + .await + .unwrap(); + let router = + SessionAffinityPushRouter::new(inner, Some(Duration::from_secs(10)), true).unwrap(); + + let mut content = request(Some(worker_id), false); + content.routing_mut().prefill_worker_id = Some(worker_id); + content.routing_mut().prefill_dp_rank = Some(0); + let observed = Arc::new(Mutex::new(None)); + let prepare_observed = observed.clone(); + let result = router + .select_and_dispatch_prefill(Context::new(content), move |request, _, dp_rank| { + request.routing_mut().prefill_dp_rank = dp_rank; + *prepare_observed.lock().unwrap() = Some(( + dp_rank, + request + .routing + .as_ref() + .and_then(|routing| routing.prefill_dp_rank), + )); + Err::<(), _>(anyhow::anyhow!("stop after preparation")) + }) + .await; + + assert!(result.is_err()); + assert_eq!(*observed.lock().unwrap(), Some((Some(0), Some(0)))); + + runtime.shutdown(); + } + #[tokio::test] async fn session_affinity_binding_wins_before_invalid_explicit_proposal() { let runtime = Runtime::from_current().unwrap(); diff --git a/lib/llm/src/session_affinity/tests.rs b/lib/llm/src/session_affinity/tests.rs index 7957109c1c8c..5b9784f9e085 100644 --- a/lib/llm/src/session_affinity/tests.rs +++ b/lib/llm/src/session_affinity/tests.rs @@ -432,15 +432,16 @@ async fn distributed_reset_and_disconnect_clear_entries() { discovery.emit(ClaimEvent::Reset); wait_for_cached_target(&coordinator, &session_id, None).await; - drop( - coordinator - .acquire(&session_id) - .await - .unwrap() - .resolve(|| target_payload(target(99, Some(0)))) - .await - .unwrap(), - ); + let resolved = coordinator + .acquire(&session_id) + .await + .unwrap() + .resolve(|| target_payload(target(99, Some(0)))) + .await + .unwrap(); + assert_eq!(resolved.target(), target(8, Some(0))); + assert!(!resolved.was_created()); + drop(resolved); discovery.disconnect(); wait_for_cached_target(&coordinator, &session_id, None).await; } diff --git a/lib/runtime/src/storage/kv/file.rs b/lib/runtime/src/storage/kv/file.rs index f4520ad48b87..8587b6c46388 100644 --- a/lib/runtime/src/storage/kv/file.rs +++ b/lib/runtime/src/storage/kv/file.rs @@ -642,10 +642,14 @@ mod tests { #[tokio::test] async fn external_delete_is_observed_under_noncanonical_root() { let t = tempfile::tempdir().unwrap(); + let canonical_root = t.path().join("canonical"); + fs::create_dir_all(&canonical_root).unwrap(); + let linked_root = t.path().join("linked"); + symlink(&canonical_root, &linked_root).unwrap(); let watcher_cancel = CancellationToken::new(); let creator_cancel = CancellationToken::new(); - let watcher_store = FileStore::new(watcher_cancel.clone(), t.path()); - let creator_store = FileStore::new(creator_cancel.clone(), t.path()); + let watcher_store = FileStore::new(watcher_cancel.clone(), &linked_root); + let creator_store = FileStore::new(creator_cancel.clone(), &canonical_root); let watcher_bucket = watcher_store .get_or_create_bucket("v1/claims", None) .await From 451d4265bbfc1109156997bcc52a86d48f337d17 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Mon, 29 Jun 2026 23:08:36 -0700 Subject: [PATCH 5/6] fix(router): propagate complete prefill targets Signed-off-by: PeaBrane --- .../src/kv_router/prefill_router/admission.rs | 4 +- lib/llm/src/kv_router/prefill_router/mod.rs | 9 +- lib/llm/src/kv_router/push_router.rs | 8 +- lib/llm/src/session_affinity/push_router.rs | 504 +++++++++++++----- lib/runtime/examples/Cargo.lock | 12 +- 5 files changed, 397 insertions(+), 140 deletions(-) diff --git a/lib/llm/src/kv_router/prefill_router/admission.rs b/lib/llm/src/kv_router/prefill_router/admission.rs index b6ecf578bed9..d4e93ab1195a 100644 --- a/lib/llm/src/kv_router/prefill_router/admission.rs +++ b/lib/llm/src/kv_router/prefill_router/admission.rs @@ -20,7 +20,7 @@ use crate::{ llm_backend::{LLMEngineOutput, PreprocessedRequest}, timing::RequestTracker, }, - session_affinity::SessionAffinityPushRouter, + session_affinity::{AffinityTarget, SessionAffinityPushRouter}, }; pub(super) enum InnerPrefillRouter { @@ -35,7 +35,7 @@ impl InnerPrefillRouter { prepare: F, ) -> Result<(M, ManyOut>)> where - F: FnOnce(&mut PreprocessedRequest, u64, Option) -> Result, + F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, { match self { InnerPrefillRouter::KvRouter(router) => { diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index cc6c42d28cad..4d68b5ec689a 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -28,6 +28,7 @@ use crate::{ preprocessor::{BootstrapInfo, PrefillResult, TraceLink}, timing::{RequestPhase, RequestTracker}, }, + session_affinity::AffinityTarget, }; mod activation; @@ -224,8 +225,8 @@ impl .ok_or_else(|| anyhow::anyhow!(PrefillError::NotActivated))?; let prefill_result: Result<(PrefillOutcome, Option)> = async { let (prepared, prefill_stream) = router - .select_and_dispatch_prefill(prefill_context, |request, worker_id, dp_rank| { - self.prepare_prefill_dispatch(request, worker_id, dp_rank) + .select_and_dispatch_prefill(prefill_context, |request, target| { + self.prepare_prefill_dispatch(request, target) }) .await?; let topology_constraints = prepared.topology_constraints; @@ -327,9 +328,9 @@ impl PrefillRouter { fn prepare_prefill_dispatch( &self, request: &mut PreprocessedRequest, - worker_id: u64, - dp_rank: Option, + target: AffinityTarget, ) -> anyhow::Result { + let AffinityTarget { worker_id, dp_rank } = target; let endpoint_id = self.endpoint_id.get(); let topology_constraints = self.preflight_kv_transfer_constraints(endpoint_id, worker_id)?; diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index 3fa261cf94a1..aa35639a2fe6 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -385,7 +385,7 @@ impl KvPushRouter { prepare: F, ) -> Result<(M, ManyOut>), Error> where - F: FnOnce(&mut PreprocessedRequest, u64, Option) -> Result, + F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, { let phase = RequestPhase::Prefill; let phase_label = phase.to_string(); @@ -398,7 +398,11 @@ impl KvPushRouter { let mut guard = self .track_selection(&request, &mut selection, is_query_only) .await?; - let metadata = match prepare(&mut request, selection.instance_id, Some(selection.dp_rank)) { + let target = AffinityTarget { + worker_id: selection.instance_id, + dp_rank: Some(selection.dp_rank), + }; + let metadata = match prepare(&mut request, target) { Ok(metadata) => metadata, Err(error) => { guard.abort().await; diff --git a/lib/llm/src/session_affinity/push_router.rs b/lib/llm/src/session_affinity/push_router.rs index 429bfb25a81c..ea44ab96b219 100644 --- a/lib/llm/src/session_affinity/push_router.rs +++ b/lib/llm/src/session_affinity/push_router.rs @@ -3,6 +3,9 @@ use std::time::Duration; +#[cfg(test)] +use std::sync::Arc; + use dynamo_runtime::{ discovery::ClaimPayloadFuture, pipeline::{ @@ -22,10 +25,15 @@ use crate::{ protocols::common::timing::{RequestPhase, WORKER_TYPE_DECODE, WORKER_TYPE_PREFILL}, }; +#[cfg(test)] +type ExactDispatchProbe = Arc; + pub struct SessionAffinityPushRouter { inner: PushRouter, affinity: Option, direct: bool, + #[cfg(test)] + exact_dispatch_probe: Option, } impl SessionAffinityPushRouter { @@ -47,6 +55,8 @@ impl SessionAffinityPushRouter { inner, affinity, direct, + #[cfg(test)] + exact_dispatch_probe: None, }) } @@ -121,13 +131,77 @@ impl SessionAffinityPushRouter { Ok((resolved, proposal_was_explicit)) } + /// Adapts the generic worker-only router while keeping a known rank attached + /// to its worker through preparation and exact dispatch. + async fn select_and_dispatch_exact_target( + &self, + request: SingleIn, + pinned_target: Option, + prepare: F, + ) -> Result<(M, ManyOut), Error> + where + F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, + { + #[cfg(test)] + let exact_dispatch_probe = self.exact_dispatch_probe.clone(); + self.inner + .select_and_dispatch_exact( + request, + pinned_target.map(|target| target.worker_id), + move |request, worker_id| { + let target = pinned_target.unwrap_or(AffinityTarget { + worker_id, + dp_rank: None, + }); + debug_assert_eq!(target.worker_id, worker_id); + let metadata = prepare(request, target)?; + #[cfg(test)] + if let Some(probe) = exact_dispatch_probe { + probe(request, target); + } + Ok(metadata) + }, + ) + .await + } + + async fn book_and_dispatch_exact_target( + &self, + request: SingleIn, + target: AffinityTarget, + advance_round_robin: bool, + prepare: F, + ) -> Result<(M, ManyOut), Error> + where + F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, + { + #[cfg(test)] + let exact_dispatch_probe = self.exact_dispatch_probe.clone(); + self.inner + .book_and_dispatch_exact( + request, + target.worker_id, + advance_round_robin, + move |request, worker_id| { + debug_assert_eq!(target.worker_id, worker_id); + let metadata = prepare(request, target)?; + #[cfg(test)] + if let Some(probe) = exact_dispatch_probe { + probe(request, target); + } + Ok(metadata) + }, + ) + .await + } + pub async fn select_and_dispatch_prefill( &self, request: SingleIn, prepare: F, ) -> Result<(M, ManyOut), Error> where - F: FnOnce(&mut PreprocessedRequest, u64, Option) -> Result, + F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, { let session_id = if self.affinity.is_some() { affinity_id(&request)? @@ -135,12 +209,9 @@ impl SessionAffinityPushRouter { None }; if !self.direct && session_id.is_none() { - let pinned_worker = phase_worker_id(&request, RequestPhase::Prefill); + let explicit = explicit_target(&request, RequestPhase::Prefill)?; return self - .inner - .select_and_dispatch_exact(request, pinned_worker, move |request, worker_id| { - prepare(request, worker_id, None) - }) + .select_and_dispatch_exact_target(request, explicit, prepare) .await; } let Some(session_id) = session_id else { @@ -150,14 +221,8 @@ impl SessionAffinityPushRouter { "worker ID required for prefill request in Direct routing mode", )); }; - let rank = target.dp_rank; return self - .inner - .select_and_dispatch_exact( - request, - Some(target.worker_id), - move |request, worker_id| prepare(request, worker_id, rank), - ) + .select_and_dispatch_exact_target(request, Some(target), prepare) .await; }; let is_query_only = request.get_annotation_value("query_instance_id").is_some(); @@ -171,21 +236,11 @@ impl SessionAffinityPushRouter { Some(target) => Some(target), None => explicit_target(&request, RequestPhase::Prefill)?, }; - let rank = selected.and_then(|target| target.dp_rank); return self - .inner - .select_and_dispatch_exact( - request, - selected.map(|target| target.worker_id), - move |request, worker_id| { - let target = AffinityTarget { - worker_id, - dp_rank: rank, - }; - Self::record_target(request, target); - prepare(request, worker_id, rank) - }, - ) + .select_and_dispatch_exact_target(request, selected, move |request, target| { + Self::record_target(request, target); + prepare(request, target) + }) .await; } @@ -200,21 +255,15 @@ impl SessionAffinityPushRouter { ) .await?; let target = resolved.target(); - let rank = target.dp_rank; let advance_round_robin = resolved.was_created() && !proposal_was_explicit; - let ((metadata, _), stream) = self - .inner - .book_and_dispatch_exact( + let (metadata, stream) = self + .book_and_dispatch_exact_target( request, - target.worker_id, + target, advance_round_robin, - move |request, worker_id| { - let target = AffinityTarget { - worker_id, - dp_rank: rank, - }; + move |request, target| { Self::record_target(request, target); - Ok((prepare(request, worker_id, rank)?, target)) + prepare(request, target) }, ) .await?; @@ -260,26 +309,14 @@ impl AsyncEngine, ManyOut, Error> Some(target) => Some(target), None => explicit_target(&request, phase)?, }; - let rank = target.and_then(|target| target.dp_rank); let (_, stream) = self - .inner - .select_and_dispatch_exact( - request, - target.map(|target| target.worker_id), - move |request, worker_id| { - if rank.is_some() { - request.routing_mut().dp_rank = rank; - } - Self::record_target( - request, - AffinityTarget { - worker_id, - dp_rank: rank, - }, - ); - Ok(()) - }, - ) + .select_and_dispatch_exact_target(request, target, move |request, target| { + if target.dp_rank.is_some() { + request.routing_mut().dp_rank = target.dp_rank; + } + Self::record_target(request, target); + Ok(()) + }) .await?; return Ok(stream); } @@ -295,24 +332,18 @@ impl AsyncEngine, ManyOut, Error> ) .await?; let target = resolved.target(); - let rank = target.dp_rank; let advance_round_robin = resolved.was_created() && !proposal_was_explicit; let (_, stream) = self - .inner - .book_and_dispatch_exact( + .book_and_dispatch_exact_target( request, - target.worker_id, + target, advance_round_robin, - move |request, worker_id| { - if rank.is_some() { - request.routing_mut().dp_rank = rank; + move |request, target| { + if target.dp_rank.is_some() { + request.routing_mut().dp_rank = target.dp_rank; } - let target = AffinityTarget { - worker_id, - dp_rank: rank, - }; Self::record_target(request, target); - Ok(target) + Ok(()) }, ) .await?; @@ -320,15 +351,6 @@ impl AsyncEngine, ManyOut, Error> } } -fn phase_worker_id(request: &PreprocessedRequest, phase: RequestPhase) -> Option { - let routing = request.routing.as_ref()?; - match phase { - RequestPhase::Prefill => routing.prefill_worker_id.or(routing.backend_instance_id), - RequestPhase::Decode => routing.decode_worker_id.or(routing.backend_instance_id), - RequestPhase::Aggregated => routing.decode_worker_id.or(routing.backend_instance_id), - } -} - #[cfg(test)] mod tests { use std::sync::{Arc, Mutex}; @@ -338,11 +360,13 @@ mod tests { distributed::DistributedConfig, pipeline::{Context, RouterMode}, }; + use futures::poll; use super::*; use crate::protocols::common::{ extensions::{SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId}, preprocessor::RoutingHints, + timing::RequestTracker, }; fn request(worker_id: Option, query_only: bool) -> PreprocessedRequest { @@ -381,6 +405,90 @@ mod tests { .expect("test router must enable affinity") } + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + struct DispatchObservation { + target: AffinityTarget, + prepared_target: AffinityTarget, + } + + struct ExactRouterHarness { + runtime: Runtime, + router: SessionAffinityPushRouter, + worker_id: u64, + observed: Arc>>, + } + + impl ExactRouterHarness { + async fn new(namespace: &str, mode: RouterMode) -> Self { + Self::new_for_phase(namespace, mode, RequestPhase::Prefill).await + } + + async fn new_for_phase(namespace: &str, mode: RouterMode, phase: RequestPhase) -> Self { + let runtime = Runtime::from_current().unwrap(); + let distributed = + DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) + .await + .unwrap(); + let endpoint = distributed + .namespace(namespace.to_string()) + .unwrap() + .component("workers".to_string()) + .unwrap() + .endpoint("prefill".to_string()); + let client = endpoint.client().await.unwrap(); + endpoint.register_endpoint_instance().await.unwrap(); + let worker_id = client.wait_for_instances().await.unwrap()[0].id(); + let observed = Arc::new(Mutex::new(Vec::new())); + let direct = mode == RouterMode::Direct; + let inner = PushRouter::from_client(client, mode).await.unwrap(); + let mut router = + SessionAffinityPushRouter::new(inner, Some(Duration::from_secs(10)), direct) + .unwrap(); + let client = router.inner.client.clone(); + let probe_observed = observed.clone(); + router.exact_dispatch_probe = Some(Arc::new(move |request, target| { + let prepared_target = explicit_target(request, phase) + .expect("prepared request must contain valid routing metadata") + .expect("prepared request must contain a routing target"); + probe_observed.lock().unwrap().push(DispatchObservation { + target, + prepared_target, + }); + client.set_overloaded_instances(&[target.worker_id]); + })); + Self { + runtime, + router, + worker_id, + observed, + } + } + + async fn dispatch_prefill(&self, request: SingleIn) { + let error = self + .router + .select_and_dispatch_prefill(request, |request, target| { + let routing = request.routing_mut(); + routing.prefill_worker_id = Some(target.worker_id); + routing.prefill_dp_rank = target.dp_rank; + Ok(target) + }) + .await + .unwrap_err(); + assert!(error.to_string().contains("overloaded")); + } + + fn observation(&self) -> DispatchObservation { + let observed = self.observed.lock().unwrap(); + assert_eq!(observed.len(), 1); + observed[0] + } + + fn shutdown(self) { + self.runtime.shutdown(); + } + } + #[tokio::test] async fn session_affinity_disabled_simple_router_has_no_coordinator() { let runtime = Runtime::from_current().unwrap(); @@ -467,7 +575,7 @@ mod tests { assert_eq!(affinity(&router).entry_count(), 0); assert!( router - .select_and_dispatch_prefill(affinity_request(None, true), |_, _, _| Ok(())) + .select_and_dispatch_prefill(affinity_request(None, true), |_, _| Ok(())) .await .is_err() ); @@ -500,7 +608,7 @@ mod tests { .contains("worker ID required for aggregated request in Direct routing mode") ); let error = router - .select_and_dispatch_prefill(Context::new(request(None, false)), |_, _, _| Ok(())) + .select_and_dispatch_prefill(Context::new(request(None, false)), |_, _| Ok(())) .await .unwrap_err(); assert!( @@ -510,61 +618,215 @@ mod tests { ); assert_eq!(affinity(&router).entry_count(), 0); - let mut decode_only = request(None, false); - decode_only.routing_mut().decode_worker_id = Some(99); + runtime.shutdown(); + } + + #[tokio::test] + async fn direct_prefill_without_session_preserves_explicit_rank_zero() { + let harness = + ExactRouterHarness::new("session_affinity_direct_prefill_rank", RouterMode::Direct) + .await; + let target = AffinityTarget { + worker_id: harness.worker_id, + dp_rank: Some(0), + }; + let mut content = request(None, false); + content.routing_mut().prefill_worker_id = Some(target.worker_id); + content.routing_mut().prefill_dp_rank = Some(0); + + harness.dispatch_prefill(Context::new(content)).await; assert_eq!( - phase_worker_id(&decode_only, RequestPhase::Aggregated), - Some(99) + harness.observation(), + DispatchObservation { + target, + prepared_target: target, + } ); + harness.shutdown(); + } - runtime.shutdown(); + #[tokio::test] + async fn non_direct_prefill_without_session_preserves_explicit_rank_zero() { + let harness = ExactRouterHarness::new( + "session_affinity_round_robin_prefill_rank", + RouterMode::RoundRobin, + ) + .await; + let target = AffinityTarget { + worker_id: harness.worker_id, + dp_rank: Some(0), + }; + let mut content = request(None, false); + content.routing_mut().prefill_worker_id = Some(target.worker_id); + content.routing_mut().prefill_dp_rank = target.dp_rank; + + harness.dispatch_prefill(Context::new(content)).await; + assert_eq!( + harness.observation(), + DispatchObservation { + target, + prepared_target: target, + } + ); + harness.shutdown(); } #[tokio::test] - async fn direct_prefill_without_session_preserves_explicit_rank_zero() { - let runtime = Runtime::from_current().unwrap(); - let distributed = - DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) + async fn distributed_affinity_target_overrides_conflicting_prefill_proposal() { + let harness = ExactRouterHarness::new( + "session_affinity_authoritative_prefill_rank", + RouterMode::RoundRobin, + ) + .await; + let authoritative = AffinityTarget { + worker_id: harness.worker_id, + dp_rank: Some(0), + }; + let seed_inner = + PushRouter::from_client(harness.router.inner.client.clone(), RouterMode::RoundRobin) .await .unwrap(); - let endpoint = distributed - .namespace("session_affinity_direct_prefill_rank".to_string()) - .unwrap() - .component("workers".to_string()) + let seed_router = + SessionAffinityPushRouter::new(seed_inner, Some(Duration::from_secs(10)), false) + .unwrap(); + let session_id = SessionAffinityId::new("authoritative-prefill-session"); + let resolved = affinity(&seed_router) + .acquire(&session_id) + .await .unwrap() - .endpoint("prefill".to_string()); - let client = endpoint.client().await.unwrap(); - endpoint.register_endpoint_instance().await.unwrap(); - let worker_id = client.wait_for_instances().await.unwrap()[0].id(); - let inner = PushRouter::from_client(client, RouterMode::Direct) + .resolve(|| Box::pin(async move { Ok(serde_json::to_value(authoritative)?) })) .await .unwrap(); - let router = - SessionAffinityPushRouter::new(inner, Some(Duration::from_secs(10)), true).unwrap(); + assert!(resolved.was_created()); + drop(resolved); + assert_eq!( + affinity(&harness.router).query_target(&session_id).unwrap(), + None + ); - let mut content = request(Some(worker_id), false); - content.routing_mut().prefill_worker_id = Some(worker_id); - content.routing_mut().prefill_dp_rank = Some(0); - let observed = Arc::new(Mutex::new(None)); - let prepare_observed = observed.clone(); - let result = router - .select_and_dispatch_prefill(Context::new(content), move |request, _, dp_rank| { - request.routing_mut().prefill_dp_rank = dp_rank; - *prepare_observed.lock().unwrap() = Some(( - dp_rank, - request - .routing - .as_ref() - .and_then(|routing| routing.prefill_dp_rank), - )); - Err::<(), _>(anyhow::anyhow!("stop after preparation")) - }) - .await; + let mut content = request(None, false); + content.routing_mut().prefill_worker_id = Some(harness.worker_id + 1); + content.routing_mut().prefill_dp_rank = Some(7); + let mut request = Context::new(content); + request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id); - assert!(result.is_err()); - assert_eq!(*observed.lock().unwrap(), Some((Some(0), Some(0)))); + harness.dispatch_prefill(request).await; + assert_eq!( + harness.observation(), + DispatchObservation { + target: authoritative, + prepared_target: authoritative, + } + ); + drop(seed_router); + harness.shutdown(); + } - runtime.shutdown(); + #[tokio::test] + async fn selected_prefill_without_rank_dispatches_with_none() { + let harness = ExactRouterHarness::new( + "session_affinity_selected_prefill_without_rank", + RouterMode::RoundRobin, + ) + .await; + let expected = AffinityTarget { + worker_id: harness.worker_id, + dp_rank: None, + }; + + harness + .dispatch_prefill(Context::new(request(None, false))) + .await; + assert_eq!( + harness.observation(), + DispatchObservation { + target: expected, + prepared_target: expected, + } + ); + harness.shutdown(); + } + + #[tokio::test] + async fn phase_barrier_keeps_prefill_and_decode_targets_separate() { + let prefill = ExactRouterHarness::new_for_phase( + "session_affinity_phase_barrier_prefill", + RouterMode::Direct, + RequestPhase::Prefill, + ) + .await; + let decode = ExactRouterHarness::new_for_phase( + "session_affinity_phase_barrier_decode", + RouterMode::Direct, + RequestPhase::Decode, + ) + .await; + let tracker = Arc::new(RequestTracker::new()); + let prefill_target = AffinityTarget { + worker_id: prefill.worker_id, + dp_rank: Some(0), + }; + let decode_target = AffinityTarget { + worker_id: decode.worker_id, + dp_rank: Some(7), + }; + assert_ne!(prefill_target.worker_id, decode_target.worker_id); + + let prefill_permit = tracker.set_phase(RequestPhase::Prefill).await; + let mut decode_transition = Box::pin(tracker.set_phase(RequestPhase::Decode)); + assert!(poll!(decode_transition.as_mut()).is_pending()); + + let mut prefill_content = request(None, false); + prefill_content.tracker = Some(tracker.clone()); + prefill_content.routing_mut().prefill_worker_id = Some(prefill_target.worker_id); + prefill_content.routing_mut().prefill_dp_rank = prefill_target.dp_rank; + let mut prefill_request = Context::new(prefill_content); + prefill_request.insert( + SESSION_AFFINITY_CONTEXT_KEY, + SessionAffinityId::new("phase-barrier-session"), + ); + prefill.dispatch_prefill(prefill_request).await; + assert_eq!( + prefill.observation(), + DispatchObservation { + target: prefill_target, + prepared_target: prefill_target, + } + ); + assert_eq!(tracker.prefill_worker_id(), Some(prefill_target.worker_id)); + assert_eq!(tracker.prefill_dp_rank(), Some(0)); + assert_eq!(tracker.decode_worker_id(), None); + assert_eq!(tracker.decode_dp_rank(), None); + + drop(prefill_permit); + let decode_permit = decode_transition.await; + + let mut decode_content = request(None, false); + decode_content.tracker = Some(tracker.clone()); + decode_content.routing_mut().decode_worker_id = Some(decode_target.worker_id); + decode_content.routing_mut().dp_rank = decode_target.dp_rank; + let mut decode_request = Context::new(decode_content); + decode_request.insert( + SESSION_AFFINITY_CONTEXT_KEY, + SessionAffinityId::new("phase-barrier-session"), + ); + let decode_error = decode.router.generate(decode_request).await.unwrap_err(); + assert!(decode_error.to_string().contains("overloaded")); + assert_eq!( + decode.observation(), + DispatchObservation { + target: decode_target, + prepared_target: decode_target, + } + ); + assert_eq!(tracker.prefill_worker_id(), Some(prefill_target.worker_id)); + assert_eq!(tracker.prefill_dp_rank(), Some(0)); + assert_eq!(tracker.decode_worker_id(), Some(decode_target.worker_id)); + assert_eq!(tracker.decode_dp_rank(), Some(7)); + + drop(decode_permit); + prefill.shutdown(); + decode.shutdown(); } #[tokio::test] diff --git a/lib/runtime/examples/Cargo.lock b/lib/runtime/examples/Cargo.lock index 2bf02373f8fd..b9e12cda7066 100644 --- a/lib/runtime/examples/Cargo.lock +++ b/lib/runtime/examples/Cargo.lock @@ -1058,15 +1058,6 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" -[[package]] -name = "fsevent-sys" -version = "4.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ee7a02da4d231650c7cea31349b889be2f45ddb3ef3032d2ec8185f6313fd2" -dependencies = [ - "libc", -] - [[package]] name = "futures" version = "0.3.32" @@ -2062,7 +2053,6 @@ checksum = "6205bd8bb1e454ad2e27422015fb5e4f2bcc7e08fa8f27058670d208324a4d2d" dependencies = [ "bitflags 2.11.1", "filetime", - "fsevent-sys", "inotify", "kqueue", "libc", @@ -2857,7 +2847,7 @@ dependencies = [ "pin-project-lite", "quinn", "rustls", - "rustls-native-certs 0.8.3", + "rustls-native-certs", "rustls-pki-types", "serde", "serde_json", From 511bede596c6fca9d23dd4e015fd826a33b18c35 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Mon, 29 Jun 2026 23:14:31 -0700 Subject: [PATCH 6/6] test(router): simplify prefill target coverage Signed-off-by: PeaBrane --- lib/llm/src/session_affinity/push_router.rs | 348 ++------------------ 1 file changed, 36 insertions(+), 312 deletions(-) diff --git a/lib/llm/src/session_affinity/push_router.rs b/lib/llm/src/session_affinity/push_router.rs index ea44ab96b219..8690c6df652c 100644 --- a/lib/llm/src/session_affinity/push_router.rs +++ b/lib/llm/src/session_affinity/push_router.rs @@ -3,9 +3,6 @@ use std::time::Duration; -#[cfg(test)] -use std::sync::Arc; - use dynamo_runtime::{ discovery::ClaimPayloadFuture, pipeline::{ @@ -25,15 +22,10 @@ use crate::{ protocols::common::timing::{RequestPhase, WORKER_TYPE_DECODE, WORKER_TYPE_PREFILL}, }; -#[cfg(test)] -type ExactDispatchProbe = Arc; - pub struct SessionAffinityPushRouter { inner: PushRouter, affinity: Option, direct: bool, - #[cfg(test)] - exact_dispatch_probe: Option, } impl SessionAffinityPushRouter { @@ -55,8 +47,6 @@ impl SessionAffinityPushRouter { inner, affinity, direct, - #[cfg(test)] - exact_dispatch_probe: None, }) } @@ -142,8 +132,6 @@ impl SessionAffinityPushRouter { where F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, { - #[cfg(test)] - let exact_dispatch_probe = self.exact_dispatch_probe.clone(); self.inner .select_and_dispatch_exact( request, @@ -154,12 +142,7 @@ impl SessionAffinityPushRouter { dp_rank: None, }); debug_assert_eq!(target.worker_id, worker_id); - let metadata = prepare(request, target)?; - #[cfg(test)] - if let Some(probe) = exact_dispatch_probe { - probe(request, target); - } - Ok(metadata) + prepare(request, target) }, ) .await @@ -175,8 +158,6 @@ impl SessionAffinityPushRouter { where F: FnOnce(&mut PreprocessedRequest, AffinityTarget) -> Result, { - #[cfg(test)] - let exact_dispatch_probe = self.exact_dispatch_probe.clone(); self.inner .book_and_dispatch_exact( request, @@ -184,12 +165,7 @@ impl SessionAffinityPushRouter { advance_round_robin, move |request, worker_id| { debug_assert_eq!(target.worker_id, worker_id); - let metadata = prepare(request, target)?; - #[cfg(test)] - if let Some(probe) = exact_dispatch_probe { - probe(request, target); - } - Ok(metadata) + prepare(request, target) }, ) .await @@ -353,20 +329,16 @@ impl AsyncEngine, ManyOut, Error> #[cfg(test)] mod tests { - use std::sync::{Arc, Mutex}; - use dynamo_runtime::{ DistributedRuntime, Runtime, distributed::DistributedConfig, pipeline::{Context, RouterMode}, }; - use futures::poll; use super::*; use crate::protocols::common::{ extensions::{SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId}, preprocessor::RoutingHints, - timing::RequestTracker, }; fn request(worker_id: Option, query_only: bool) -> PreprocessedRequest { @@ -405,90 +377,6 @@ mod tests { .expect("test router must enable affinity") } - #[derive(Clone, Copy, Debug, PartialEq, Eq)] - struct DispatchObservation { - target: AffinityTarget, - prepared_target: AffinityTarget, - } - - struct ExactRouterHarness { - runtime: Runtime, - router: SessionAffinityPushRouter, - worker_id: u64, - observed: Arc>>, - } - - impl ExactRouterHarness { - async fn new(namespace: &str, mode: RouterMode) -> Self { - Self::new_for_phase(namespace, mode, RequestPhase::Prefill).await - } - - async fn new_for_phase(namespace: &str, mode: RouterMode, phase: RequestPhase) -> Self { - let runtime = Runtime::from_current().unwrap(); - let distributed = - DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) - .await - .unwrap(); - let endpoint = distributed - .namespace(namespace.to_string()) - .unwrap() - .component("workers".to_string()) - .unwrap() - .endpoint("prefill".to_string()); - let client = endpoint.client().await.unwrap(); - endpoint.register_endpoint_instance().await.unwrap(); - let worker_id = client.wait_for_instances().await.unwrap()[0].id(); - let observed = Arc::new(Mutex::new(Vec::new())); - let direct = mode == RouterMode::Direct; - let inner = PushRouter::from_client(client, mode).await.unwrap(); - let mut router = - SessionAffinityPushRouter::new(inner, Some(Duration::from_secs(10)), direct) - .unwrap(); - let client = router.inner.client.clone(); - let probe_observed = observed.clone(); - router.exact_dispatch_probe = Some(Arc::new(move |request, target| { - let prepared_target = explicit_target(request, phase) - .expect("prepared request must contain valid routing metadata") - .expect("prepared request must contain a routing target"); - probe_observed.lock().unwrap().push(DispatchObservation { - target, - prepared_target, - }); - client.set_overloaded_instances(&[target.worker_id]); - })); - Self { - runtime, - router, - worker_id, - observed, - } - } - - async fn dispatch_prefill(&self, request: SingleIn) { - let error = self - .router - .select_and_dispatch_prefill(request, |request, target| { - let routing = request.routing_mut(); - routing.prefill_worker_id = Some(target.worker_id); - routing.prefill_dp_rank = target.dp_rank; - Ok(target) - }) - .await - .unwrap_err(); - assert!(error.to_string().contains("overloaded")); - } - - fn observation(&self) -> DispatchObservation { - let observed = self.observed.lock().unwrap(); - assert_eq!(observed.len(), 1); - observed[0] - } - - fn shutdown(self) { - self.runtime.shutdown(); - } - } - #[tokio::test] async fn session_affinity_disabled_simple_router_has_no_coordinator() { let runtime = Runtime::from_current().unwrap(); @@ -622,211 +510,47 @@ mod tests { } #[tokio::test] - async fn direct_prefill_without_session_preserves_explicit_rank_zero() { - let harness = - ExactRouterHarness::new("session_affinity_direct_prefill_rank", RouterMode::Direct) - .await; - let target = AffinityTarget { - worker_id: harness.worker_id, - dp_rank: Some(0), - }; - let mut content = request(None, false); - content.routing_mut().prefill_worker_id = Some(target.worker_id); - content.routing_mut().prefill_dp_rank = Some(0); - - harness.dispatch_prefill(Context::new(content)).await; - assert_eq!( - harness.observation(), - DispatchObservation { - target, - prepared_target: target, - } - ); - harness.shutdown(); - } - - #[tokio::test] - async fn non_direct_prefill_without_session_preserves_explicit_rank_zero() { - let harness = ExactRouterHarness::new( - "session_affinity_round_robin_prefill_rank", - RouterMode::RoundRobin, - ) - .await; - let target = AffinityTarget { - worker_id: harness.worker_id, - dp_rank: Some(0), - }; - let mut content = request(None, false); - content.routing_mut().prefill_worker_id = Some(target.worker_id); - content.routing_mut().prefill_dp_rank = target.dp_rank; - - harness.dispatch_prefill(Context::new(content)).await; - assert_eq!( - harness.observation(), - DispatchObservation { - target, - prepared_target: target, - } - ); - harness.shutdown(); - } - - #[tokio::test] - async fn distributed_affinity_target_overrides_conflicting_prefill_proposal() { - let harness = ExactRouterHarness::new( - "session_affinity_authoritative_prefill_rank", - RouterMode::RoundRobin, - ) - .await; - let authoritative = AffinityTarget { - worker_id: harness.worker_id, - dp_rank: Some(0), - }; - let seed_inner = - PushRouter::from_client(harness.router.inner.client.clone(), RouterMode::RoundRobin) + async fn prefill_preparation_receives_explicit_rank_zero() { + let runtime = Runtime::from_current().unwrap(); + let distributed = + DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local()) .await .unwrap(); - let seed_router = - SessionAffinityPushRouter::new(seed_inner, Some(Duration::from_secs(10)), false) - .unwrap(); - let session_id = SessionAffinityId::new("authoritative-prefill-session"); - let resolved = affinity(&seed_router) - .acquire(&session_id) - .await + let endpoint = distributed + .namespace("session_affinity_prefill_target".to_string()) .unwrap() - .resolve(|| Box::pin(async move { Ok(serde_json::to_value(authoritative)?) })) - .await - .unwrap(); - assert!(resolved.was_created()); - drop(resolved); - assert_eq!( - affinity(&harness.router).query_target(&session_id).unwrap(), - None - ); - - let mut content = request(None, false); - content.routing_mut().prefill_worker_id = Some(harness.worker_id + 1); - content.routing_mut().prefill_dp_rank = Some(7); - let mut request = Context::new(content); - request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id); - - harness.dispatch_prefill(request).await; - assert_eq!( - harness.observation(), - DispatchObservation { - target: authoritative, - prepared_target: authoritative, - } - ); - drop(seed_router); - harness.shutdown(); - } - - #[tokio::test] - async fn selected_prefill_without_rank_dispatches_with_none() { - let harness = ExactRouterHarness::new( - "session_affinity_selected_prefill_without_rank", - RouterMode::RoundRobin, - ) - .await; + .component("workers".to_string()) + .unwrap() + .endpoint("prefill".to_string()); + let client = endpoint.client().await.unwrap(); + endpoint.register_endpoint_instance().await.unwrap(); + let worker_id = client.wait_for_instances().await.unwrap()[0].id(); let expected = AffinityTarget { - worker_id: harness.worker_id, - dp_rank: None, + worker_id, + dp_rank: Some(0), }; - harness - .dispatch_prefill(Context::new(request(None, false))) - .await; - assert_eq!( - harness.observation(), - DispatchObservation { - target: expected, - prepared_target: expected, - } - ); - harness.shutdown(); - } + for (mode, direct) in [(RouterMode::Direct, true), (RouterMode::RoundRobin, false)] { + let inner = PushRouter::from_client(client.clone(), mode).await.unwrap(); + let router = SessionAffinityPushRouter::new(inner, None, direct).unwrap(); + let mut content = request(None, false); + content.routing_mut().prefill_worker_id = Some(worker_id); + content.routing_mut().prefill_dp_rank = Some(0); + let mut observed = None; + + let error = router + .select_and_dispatch_prefill(Context::new(content), |_, target| { + observed = Some(target); + Err::<(), _>(anyhow::anyhow!("stop before dispatch")) + }) + .await + .unwrap_err(); - #[tokio::test] - async fn phase_barrier_keeps_prefill_and_decode_targets_separate() { - let prefill = ExactRouterHarness::new_for_phase( - "session_affinity_phase_barrier_prefill", - RouterMode::Direct, - RequestPhase::Prefill, - ) - .await; - let decode = ExactRouterHarness::new_for_phase( - "session_affinity_phase_barrier_decode", - RouterMode::Direct, - RequestPhase::Decode, - ) - .await; - let tracker = Arc::new(RequestTracker::new()); - let prefill_target = AffinityTarget { - worker_id: prefill.worker_id, - dp_rank: Some(0), - }; - let decode_target = AffinityTarget { - worker_id: decode.worker_id, - dp_rank: Some(7), - }; - assert_ne!(prefill_target.worker_id, decode_target.worker_id); - - let prefill_permit = tracker.set_phase(RequestPhase::Prefill).await; - let mut decode_transition = Box::pin(tracker.set_phase(RequestPhase::Decode)); - assert!(poll!(decode_transition.as_mut()).is_pending()); - - let mut prefill_content = request(None, false); - prefill_content.tracker = Some(tracker.clone()); - prefill_content.routing_mut().prefill_worker_id = Some(prefill_target.worker_id); - prefill_content.routing_mut().prefill_dp_rank = prefill_target.dp_rank; - let mut prefill_request = Context::new(prefill_content); - prefill_request.insert( - SESSION_AFFINITY_CONTEXT_KEY, - SessionAffinityId::new("phase-barrier-session"), - ); - prefill.dispatch_prefill(prefill_request).await; - assert_eq!( - prefill.observation(), - DispatchObservation { - target: prefill_target, - prepared_target: prefill_target, - } - ); - assert_eq!(tracker.prefill_worker_id(), Some(prefill_target.worker_id)); - assert_eq!(tracker.prefill_dp_rank(), Some(0)); - assert_eq!(tracker.decode_worker_id(), None); - assert_eq!(tracker.decode_dp_rank(), None); - - drop(prefill_permit); - let decode_permit = decode_transition.await; - - let mut decode_content = request(None, false); - decode_content.tracker = Some(tracker.clone()); - decode_content.routing_mut().decode_worker_id = Some(decode_target.worker_id); - decode_content.routing_mut().dp_rank = decode_target.dp_rank; - let mut decode_request = Context::new(decode_content); - decode_request.insert( - SESSION_AFFINITY_CONTEXT_KEY, - SessionAffinityId::new("phase-barrier-session"), - ); - let decode_error = decode.router.generate(decode_request).await.unwrap_err(); - assert!(decode_error.to_string().contains("overloaded")); - assert_eq!( - decode.observation(), - DispatchObservation { - target: decode_target, - prepared_target: decode_target, - } - ); - assert_eq!(tracker.prefill_worker_id(), Some(prefill_target.worker_id)); - assert_eq!(tracker.prefill_dp_rank(), Some(0)); - assert_eq!(tracker.decode_worker_id(), Some(decode_target.worker_id)); - assert_eq!(tracker.decode_dp_rank(), Some(7)); - - drop(decode_permit); - prefill.shutdown(); - decode.shutdown(); + assert!(error.to_string().contains("stop before dispatch")); + assert_eq!(observed, Some(expected)); + } + + runtime.shutdown(); } #[tokio::test]