From be3de2170d9c791f4f042c664f4c706c9f75223f Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Tue, 6 Oct 2026 02:19:23 -0700 Subject: [PATCH] feat(kv-router): one session affinity config and validator, mode on every host `SessionAffinityConfig` is the one session-affinity configuration and gains the one validator: `validate` checks the TTL range and the limits before a table starts, and `ttl_from_secs_f64` is the one range check for a TTL given in seconds. The four copies of the TTL check (the table, the selection service builder, the Python bindings, and the frontend's argument parser) and the EPP's separate range attribute all go through it. The bindings export `validate_session_affinity_ttl_secs` so `dynamo.frontend`'s argument parser reports its own flag with the router's check instead of a copied constant. The session-affinity mode (`hard` or `soft`, `hard` by default) is now settable on the standalone selection service (`--session-affinity-mode`, `SelectionService(session_affinity_mode=...)`) and the EPP (`DYN_EPP_SESSION_AFFINITY_MODE`), which were always `hard`. Docs: the frontend's mode table describes the pin rule (a hard binding limits selection while eligible, re-binds when its worker leaves, and is rejected before dispatch when the request's own constraints exclude it; a full table routes without affinity), the standalone service and EPP reference pages gain the mode setting, the custom-policy guide describes hard affinity as a narrowed candidate set rather than an exact-target path, and the replication subjects and topics are documented with their one-release dual-publish window. Signed-off-by: Ishan Dhanani --- .../src/dynamo/frontend/frontend_args.py | 12 +-- .../ext-proc/src/epp_standalone_config.rs | 17 ++-- .../ext-proc/src/selector.rs | 12 ++- .../ext-proc/src/topology_adapter.rs | 1 + .../router/configuration-and-tuning.md | 16 ++- .../router/custom-worker-selection.mdx | 2 +- .../router/standalone-selection.md | 5 +- .../kv-aware-routing/vanilla-vllm-onramp.mdx | 1 + .../components/frontend-configuration.mdx | 8 ++ lib/bindings/python/rust/lib.rs | 4 + lib/bindings/python/rust/llm/kv.rs | 50 +++++++--- lib/bindings/python/src/dynamo/_core.pyi | 14 ++- .../src/services/selection/affinity/README.md | 4 + .../src/services/selection/affinity/table.rs | 97 ++++++++++++++++--- .../src/services/selection/core/mod.rs | 2 + .../src/services/selection/service.rs | 24 ++++- lib/kv-router/src/services/selection/tests.rs | 1 + lib/llm/src/session_affinity/mod.rs | 3 +- 18 files changed, 217 insertions(+), 56 deletions(-) diff --git a/components/src/dynamo/frontend/frontend_args.py b/components/src/dynamo/frontend/frontend_args.py index 020c41e195ee..7f556d184ac6 100644 --- a/components/src/dynamo/frontend/frontend_args.py +++ b/components/src/dynamo/frontend/frontend_args.py @@ -21,6 +21,7 @@ RouterArgGroup, RouterConfigBase, ) +from dynamo._core import validate_session_affinity_ttl_secs from dynamo.common.configuration.utils import ( add_argument, add_negatable_bool_argument, @@ -31,7 +32,6 @@ from . import __version__ _U32_MAX = 2**32 - 1 -_MAX_SESSION_AFFINITY_TTL_SECS = 31_536_000 def validate_model_name(value: str) -> str: @@ -141,12 +141,10 @@ def validate(self) -> None: ) if self.min_initial_workers < 0: raise ValueError("--router-min-initial-workers must be >= 0") - if self.session_affinity_ttl_secs is not None and not ( - 1 <= self.session_affinity_ttl_secs <= _MAX_SESSION_AFFINITY_TTL_SECS - ): - raise ValueError( - "--router-session-affinity-ttl-secs must be between 1 and " - f"{_MAX_SESSION_AFFINITY_TTL_SECS}" + if self.session_affinity_ttl_secs is not None: + # The same range check the router applies, naming this flag. + validate_session_affinity_ttl_secs( + self.session_affinity_ttl_secs, "--router-session-affinity-ttl-secs" ) if self.tokenizer_backend not in self._VALID_TOKENIZER_BACKENDS: raise ValueError( diff --git a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs index d8a6b38610fb..c828cb9f04f1 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_standalone_config.rs @@ -10,6 +10,7 @@ //! [`EppStandaloneConfig::from_env`] reads envs, applies defaults, and calls //! [`EppStandaloneConfig::validate_config`] for field and cross-field checks. +use dynamo_kv_router::services::selection::affinity::SessionAffinityMode; use validator::Validate; use validator::ValidationError; @@ -176,13 +177,12 @@ pub struct EppStandaloneConfig { pub max_inflight_requests: usize, /// Pin each `x-dynamo-session-id` to the worker that served it for this /// long after its last request (`DYN_EPP_SESSION_AFFINITY_TTL_SECS`). - /// `None` disables session affinity. - #[validate(range( - min = 1.0, - max = 31536000.0, - message = "DYN_EPP_SESSION_AFFINITY_TTL_SECS must be between 1 and 31536000 seconds" - ))] + /// `None` disables session affinity. The range is checked by the shared + /// `SessionAffinityConfig` validator when the selector starts. pub session_affinity_ttl_secs: Option, + /// How a bound session treats a dispatch that landed elsewhere + /// (`DYN_EPP_SESSION_AFFINITY_MODE`, `hard` or `soft`); `hard` unless set. + pub session_affinity_mode: SessionAffinityMode, } impl EppStandaloneConfig { @@ -247,6 +247,11 @@ impl EppStandaloneConfig { max_inflight_requests: opt_parse::(get, "DYN_EPP_MAX_INFLIGHT_REQUESTS")? .unwrap_or(DEFAULT_MAX_INFLIGHT_REQUESTS), session_affinity_ttl_secs: opt_parse::(get, "DYN_EPP_SESSION_AFFINITY_TTL_SECS")?, + session_affinity_mode: opt_parse::( + get, + "DYN_EPP_SESSION_AFFINITY_MODE", + )? + .unwrap_or_default(), }) } diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 790ea6d1510f..7f26c55ac368 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -14,6 +14,7 @@ use anyhow::{Context, Result, anyhow}; use dynamo_kv_router::config::{KvRouterConfig, try_kv_router_config_from_dynamo_env}; use dynamo_kv_router::protocols::RoutingConstraints; +use dynamo_kv_router::services::selection::affinity::SessionAffinityConfig; use dynamo_kv_router::services::selection::{ PromptRequest, SelectAndReserveRequest as CoreSelectAndReserveRequest, SelectionError, SelectionService, SelectionServiceBuilder, WorkerSelectionPolicyRegistry, @@ -120,10 +121,12 @@ impl Selector { builder = builder.replica_sync(peer_replication.sync_port, Vec::new()); } if let Some(ttl) = cfg.session_affinity_ttl_secs { - builder = builder.session_affinity( - std::time::Duration::try_from_secs_f64(ttl) - .context("invalid session affinity TTL")?, - ); + builder = builder + .session_affinity( + SessionAffinityConfig::ttl_from_secs_f64(ttl) + .context("invalid session affinity TTL")?, + ) + .session_affinity_mode(cfg.session_affinity_mode); } let service = Arc::new( builder @@ -363,6 +366,7 @@ models: max_num_batched_tokens: Some(8192), max_inflight_requests: 1024, session_affinity_ttl_secs: None, + session_affinity_mode: Default::default(), } } diff --git a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs index 244381551259..7cde762f5a34 100644 --- a/deploy/inference-gateway/ext-proc/src/topology_adapter.rs +++ b/deploy/inference-gateway/ext-proc/src/topology_adapter.rs @@ -153,6 +153,7 @@ mod tests { max_num_batched_tokens: Some(8192), max_inflight_requests: 1024, session_affinity_ttl_secs: None, + session_affinity_mode: Default::default(), } } diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md index cff26c43df95..295135a178cc 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md @@ -336,10 +336,12 @@ The first successfully dispatched request binds the session ID to its selected w | Mode | Behavior | |---|---| -| `hard` | Default. Exact-dispatch to the stored target. If the worker or rank is no longer valid, invalidate the binding and retry normal selection once | -| `soft` | Pass the stored target through the normal selection pipeline as an advisory target. The built-in selector retains it while eligible; a custom policy can choose another worker | +| `hard` | Default. Selection is limited to the bound worker (and rank) while it is eligible, for every selection policy. A hard-bound request goes to its worker and queues there; the router holds it only when every eligible worker is busy, as for any request. A bound worker that discovery no longer knows drops the binding before selection and the session re-binds to the worker the request runs on. A bound worker that is live but excluded by the request's own constraints (allowed workers, taints, migration exclusions, or the overload detector) selects another worker and the request is rejected before dispatch as invalid, keeping the binding | +| `soft` | The stored target enters the normal selection pipeline as a preference. The built-in selector keeps it while it is a candidate; a custom policy can choose another worker, and the session rebinds to the worker the request ran on | -**Experimental.** Available since v1.6. When session affinity is enabled (`--router-session-affinity-ttl-secs`), a request that carries a parent session id binds under an internal key derived from that parent id instead of its own session, so the subagents of one parent share a binding while the parent keeps its own. Dynamo's affinity coordinator owns that binding: it commits only after a successful dispatch, is version-checked against concurrent updates, expires on the same TTL, and counts against the same global entry limit as any session binding. A request that carries an explicit worker target stays on its own session, so it is neither rejected against the group nor able to move it. Under the default `hard` mode the group is pinned to its first worker and does not migrate because of load; the binding resets only when the target becomes unusable or dispatch fails. +Explicit request targets (`nvext.backend_instance_id`, `decode_worker_id`, `prefill_worker_id`) stay exact in both modes and must agree with an existing binding. When the affinity table is full (65,536 sessions per router), a new session routes without affinity rather than failing; the router counts these fallbacks. + +**Experimental.** Available since v1.6. When session affinity is enabled (`--router-session-affinity-ttl-secs`), a request that carries a parent session id binds under an internal key derived from that parent id instead of its own session, so the subagents of one parent share a binding while the parent keeps its own. Dynamo's session resolver owns that binding: it commits only after a successful dispatch, is version-checked against concurrent updates, expires on the same TTL, and counts against the same global entry limit as any session binding. A request that carries an explicit worker target stays on its own session, so it is neither rejected against the group nor able to move it. Under the default `hard` mode the group is pinned to its first worker and does not migrate because of load; the binding resets only when the bound worker leaves the pool. One behavior is known and unresolved: because siblings share one binding, a concurrent fan-out waits for the first sibling's dispatch to commit before the others are placed. @@ -354,7 +356,7 @@ python -m dynamo.frontend \ --router-session-affinity-mode soft ``` -Concurrent requests can share a binding. Versioned updates prevent an older concurrent request from replacing a newer soft rebind. Active requests prevent expiry. When a request lease ends after EOF, early drop, error, or cancellation, the idle timer restarts. A missing hard-bound worker or a non-cancellation hard-mode selection, setup, dispatch, or target-validation failure invalidates the binding. +Concurrent requests can share a binding. Versioned updates prevent an older concurrent request from replacing a newer soft rebind. Active requests prevent expiry. When a request lease ends after EOF, early drop, error, or cancellation, the idle timer restarts. A hard-bound worker that leaves the pool invalidates the binding before the next request selects; a failed dispatch binds nothing and leaves an existing binding as it was. The same session layer serves the frontend (KV and builtin routing modes), the standalone selection service, and the EPP, so a binding means the same thing on every host. The configured value is the idle timeout. It is independent of `--router-ttl-secs` and `--router-predicted-ttl-secs`. Omit the session-affinity @@ -368,7 +370,11 @@ When session affinity is enabled, routers synchronize affinity bindings through Runtime event plane. The origin publishes a binding after successful dispatch so concurrent requests can observe it, then publishes it again when the request lease ends so peer idle timers restart when the request becomes idle. The extra event -fanout is an intentional tradeoff. +fanout is an intentional tradeoff. Bindings travel on the `session_affinity_events_v2` +subject, whose payload names the routing partition (model name and routing group) +so a binding applies only to the partition it belongs to; for one release routers +also publish and apply the older partition-less `session_affinity_events` subject so +mixed-version replicas keep converging. Synchronization is advisory. Each replica owns its local idle TTL and uses the first live binding it observes. A matching update refreshes that local deadline, diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 15e97fd1639f..f4f69bbadc86 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -415,7 +415,7 @@ flowchart LR Match both fields when a policy wants to recognize the target. A target with `dp_rank: None` matches every rank on its worker; a populated rank matches only that worker-rank pair. The target can be absent from the candidate table when the worker is unavailable or a custom filter rejects it. -Hard affinity remains the default. A hard binding and every explicit request target take the exact-target path instead of advisory custom selection. A selection, setup, or dispatch failure before a response stream leaves the previous soft binding intact. An error or cancellation after the stream is returned does not roll back a completed rebind. +Hard affinity remains the default. Under `hard`, the host limits the candidate set to the bound worker while that worker is eligible, for every policy: your filters, scorers, and picker still run, but only over that worker's data-parallel ranks, and `affinity_target()` names the binding they are already limited to. Explicit request targets are exact in both modes. A selection, setup, or dispatch failure before a response stream leaves the previous binding intact. An error or cancellation after the stream is returned does not roll back a completed rebind. The [custom policy examples](https://github.com/ai-dynamo/dynamo/blob/main/examples/router/custom-policy-example/README.md) include an overload-aware soft pinning policy that retains the advisory target until its active-request count exceeds a configured threshold, then selects the least-loaded alternative. Its two-Mocker walkthrough holds the first request open and verifies an `A -> B -> B` worker sequence: overload moves the second request to B, and the third request retains B after load drains, proving that the plugin-selected dispatch updated the binding. diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index 186ae71bab5d..37672ab8246a 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -68,7 +68,8 @@ APIs. Those bindings should wrap `SelectionService` rather than construct | `--indexer-peers` | none | Comma-separated HTTP URLs used for startup KV recovery through `/dump`. | | `--replica-sync-port` | none | Local ZMQ PUB port for active-load lifecycle events. The selector binds `tcp://*:` internally. | | `--replica-sync-peers` | none | Comma-separated ZMQ PUB endpoints for selector peers. Requires `--replica-sync-port`. | -| `--session-affinity-ttl-secs` | none | Pin each request `session_id` to the worker that served it for this long after its last request (1 to 31,536,000 seconds). When `--replica-sync-port` is configured, bindings replicate to peers over the replica mesh (`dynamo.session-affinity.v1` topic). | +| `--session-affinity-ttl-secs` | none | Pin each request `session_id` to the worker that served it for this long after its last request (1 to 31,536,000 seconds). When `--replica-sync-port` is configured, bindings replicate to peers over the replica mesh (`dynamo.session-affinity.v2` topic, with `.v1` published and applied alongside it for one release). | +| `--session-affinity-mode` | `hard` | How a bound session treats a dispatch that landed elsewhere. `hard` limits selection to the bound worker while it is eligible and rejects a booking elsewhere; `soft` offers the binding to the selection policy as a preference and rebinds to the worker booked. | | `--selection-cache-ttl-secs` | `120` | Seconds an unclaimed pending selection lives before eviction. | | `--selection-cache-max-entries` | `4096` | Maximum resident pending selections, evicting oldest first. | | `--selection-cache-max-bytes` | `268435456` | Approximate byte budget across resident pending selections. | @@ -346,7 +347,7 @@ hash producer with the same algorithm, key, and key ID as the selector. ### `session_id` -Both `POST /select` and `POST /select_and_reserve` accept an optional `session_id` string. With `--session-affinity-ttl-secs` enabled, each model and routing group has its own binding table. `/select_and_reserve` binds a new session to the selected worker and holds its lease until the reservation is released or expires. The idle TTL starts when the last lease is released. `/select` can use an existing binding but does not create one. +Both `POST /select` and `POST /select_and_reserve` accept an optional `session_id` string. With `--session-affinity-ttl-secs` enabled, each model and routing group has its own binding table. `/select_and_reserve` binds a new session to the selected worker and holds its lease until the reservation is released or expires. The idle TTL starts when the last lease is released. `/select` can use an existing binding but does not create one. Under `--session-affinity-mode hard` (the default) selection is limited to the bound worker while it is eligible, for every policy; a bound worker that is no longer schedulable drops the binding and the session re-binds to the worker booked. Under `soft` the binding is a preference the policy may override. Bindings replicate over the configured replica mesh. An explicit `affinity_target` or `pinned_worker` takes precedence over the session binding. Without a configured TTL, `session_id` is only policy input: custom policies can read it through `WorkerSelectionContext::session_id()`, while the built-in selector ignores it. See [Write Custom Routing Strategies](custom-worker-selection.mdx). diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index ec91a4bc9b05..2b5be7d8c358 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx +++ b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx @@ -510,6 +510,7 @@ For a single EPP replica, set `spec.replicas: 1` and remove `DYN_EPP_PEER_SERVIC | `DYN_EPP_TOKENIZATION_TIMEOUT_MS` | No | Render request timeout; defaults to `5000`. | | `DYN_EPP_TOKENIZER_MAX_RESPONSE_BYTES` | No | Maximum render response size; defaults to `16777216`. | | `DYN_EPP_SESSION_AFFINITY_TTL_SECS` | No | Pin each `x-dynamo-session-id` to the worker that served it for this many seconds (1 to 31,536,000) after its last request; the binding replicates to peer EPPs over the replica mesh. Unset disables session affinity. | +| `DYN_EPP_SESSION_AFFINITY_MODE` | No | `hard` (default) keeps a bound session on its worker while that worker is eligible; `soft` offers the binding to the selection policy as a preference and rebinds to the worker booked. | | `DYN_EPP_MAX_INFLIGHT_REQUESTS` | No | Concurrent EPP request guardrail; defaults to `1024`, and excess requests receive `503`. | | `DYN_EPP_PEER_SERVICE` | No | Existing EPP Service used to discover sibling replicas. Enables replica sync. | | `DYN_EPP_REPLICA_SYNC_PORT` | No | ZMQ port every replica binds and dials for replica sync; defaults to `9092`. All replicas selected by `DYN_EPP_PEER_SERVICE` must use the same value. | diff --git a/docs/fern/pages/reference/components/frontend-configuration.mdx b/docs/fern/pages/reference/components/frontend-configuration.mdx index 635fa19c0b86..33fd5fc66519 100644 --- a/docs/fern/pages/reference/components/frontend-configuration.mdx +++ b/docs/fern/pages/reference/components/frontend-configuration.mdx @@ -127,6 +127,14 @@ router. See [Router Guide](../../developer-guide/knowledge-base/modular-componen Environment variable: `DYN_ROUTER_SESSION_AFFINITY_TTL_SECS` + + How an existing session binding takes part in worker selection. `hard` limits selection to the + bound worker while it is eligible; `soft` offers the binding to the selection policy as a + preference. See [Session affinity](../../developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md#session-affinity). + + Environment variable: `DYN_ROUTER_SESSION_AFFINITY_MODE` + + Fall back to aggregated mode when prefill workers are unavailable. diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index d455c0e6c83e..2eb806432230 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -317,6 +317,10 @@ fn register_core(m: &Bound<'_, PyModule>) -> PyResult<()> { )?; m.add_function(wrap_pyfunction!(llm::kv::compute_block_hash_for_seq_py, m)?)?; + m.add_function(wrap_pyfunction!( + llm::kv::validate_session_affinity_ttl_secs_py, + m + )?)?; m.add_function(wrap_pyfunction!(lora_name_to_id, m)?)?; #[cfg(feature = "mm-routing")] m.add_function(wrap_pyfunction!(resolve_routing_image_token_id, m)?)?; diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 81946cabd7f6..afa2f6e2f4cd 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -53,7 +53,8 @@ use llm_rs::kv_router::publisher::{KvEventSourceConfig, create_stored_blocks}; use llm_rs::protocols::common::timing::RequestTracker; use llm_rs::protocols::common::{OutputOptions, SamplingOptions, StopConditions}; use llm_rs::session_affinity::{ - MAX_SESSION_AFFINITY_TTL_SECS, SessionAffinityMode as RsSessionAffinityMode, + MAX_SESSION_AFFINITY_TTL_SECS, MIN_SESSION_AFFINITY_TTL_SECS, + SessionAffinityConfig as RsSessionAffinityConfig, SessionAffinityMode as RsSessionAffinityMode, }; use super::ais_callback::create_ais_prefill_load_estimator; @@ -371,6 +372,12 @@ struct SelectServiceCli { #[arg(long)] session_affinity_ttl_secs: Option, + /// How a bound session treats a dispatch that landed elsewhere: `hard` + /// keeps every request on the bound worker while it is eligible, `soft` + /// offers the binding to the selection policy as a preference + #[arg(long, default_value = "hard")] + session_affinity_mode: RsSessionAffinityMode, + /// Seconds an unclaimed pending selection lives before eviction #[arg(long)] selection_cache_ttl_secs: Option, @@ -558,6 +565,7 @@ where .map(session_affinity_ttl_from_secs) .transpose() .map_err(anyhow::Error::msg)?, + session_affinity_mode: cli.session_affinity_mode, kv_router_config, selection_cache: selection_cache_config_from_overrides( cli.selection_cache_ttl_secs, @@ -634,13 +642,17 @@ impl SelectionCacheConfig { } } +/// The one TTL range check (`SessionAffinityConfig::ttl_from_secs_f64`), with +/// the Python-facing parameter named in the error. +fn session_affinity_ttl_range_error(name: &str) -> String { + format!( + "{name} must be between {MIN_SESSION_AFFINITY_TTL_SECS} and {MAX_SESSION_AFFINITY_TTL_SECS}" + ) +} + fn session_affinity_ttl_from_secs(ttl: f64) -> Result { - if !(1.0..=MAX_SESSION_AFFINITY_TTL_SECS as f64).contains(&ttl) { - return Err(format!( - "session_affinity_ttl_secs must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS}" - )); - } - Ok(Duration::from_secs_f64(ttl)) + RsSessionAffinityConfig::ttl_from_secs_f64(ttl) + .map_err(|_| session_affinity_ttl_range_error("session_affinity_ttl_secs")) } /// Range check for a whole-second TTL; `None` passes. @@ -650,6 +662,16 @@ pub(crate) fn check_session_affinity_ttl_secs(ttl: Option) -> PyResult<()> Ok(()) } +/// Check a session-affinity idle TTL in seconds against the supported range, +/// naming `name` in the error so an argument parser reports its own flag. +#[pyfunction] +#[pyo3(name = "validate_session_affinity_ttl_secs", signature = (ttl_secs, name = "session_affinity_ttl_secs"))] +pub fn validate_session_affinity_ttl_secs_py(ttl_secs: f64, name: &str) -> PyResult<()> { + RsSessionAffinityConfig::ttl_from_secs_f64(ttl_secs) + .map(|_| ()) + .map_err(|_| PyValueError::new_err(session_affinity_ttl_range_error(name))) +} + /// In-process handle to a managed Dynamo `SelectionService`. #[cfg(feature = "select-service")] #[pyclass] @@ -662,7 +684,7 @@ pub(crate) struct SelectionService { impl SelectionService { /// Create a selection service. `indexer_threads` sizes the KV indexer pool. #[new] - #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, session_affinity_ttl_secs = None))] + #[pyo3(signature = (*, indexer_threads = 4, indexer_peers = None, replica_sync_port = None, replica_sync_peers = None, selection_cache = None, session_affinity_ttl_secs = None, session_affinity_mode = "hard"))] #[allow(clippy::too_many_arguments)] fn new( py: Python<'_>, @@ -672,7 +694,11 @@ impl SelectionService { replica_sync_peers: Option>, selection_cache: Option, session_affinity_ttl_secs: Option, + session_affinity_mode: &str, ) -> PyResult { + let session_affinity_mode = session_affinity_mode + .parse::() + .map_err(PyValueError::new_err)?; let replica_sync_peers = replica_sync_peers.unwrap_or_default(); if replica_sync_port.is_none() && !replica_sync_peers.is_empty() { return Err(PyValueError::new_err( @@ -695,9 +721,11 @@ impl SelectionService { builder = builder.replica_sync(port, replica_sync_peers); } if let Some(ttl) = session_affinity_ttl_secs { - builder = builder.session_affinity( - session_affinity_ttl_from_secs(ttl).map_err(PyValueError::new_err)?, - ); + builder = builder + .session_affinity( + session_affinity_ttl_from_secs(ttl).map_err(PyValueError::new_err)?, + ) + .session_affinity_mode(session_affinity_mode); } let inner = py .allow_threads(|| crate::bridge_runtime().block_on(builder.build())) diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index c14b7629feea..0e25437e4408 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -413,6 +413,15 @@ class ModelCardInstanceId: ... +def validate_session_affinity_ttl_secs( + ttl_secs: float, name: str = "session_affinity_ttl_secs" +) -> None: + """Raise ``ValueError`` unless ``ttl_secs`` is a finite session-affinity idle + TTL between 1 and 31536000 seconds. ``name`` is the parameter to blame in + the error, so an argument parser can report its own flag. + """ + ... + def compute_block_hash_for_seq( tokens: List[int], kv_block_size: int, @@ -765,11 +774,14 @@ class SelectionService: replica_sync_peers: Optional[list[str]] = None, selection_cache: Optional[SelectionCacheConfig] = None, session_affinity_ttl_secs: Optional[float] = None, + session_affinity_mode: str = "hard", ) -> None: """Create a selection service. `indexer_threads` sizes the KV indexer pool. `session_affinity_ttl_secs` enables session affinity with an idle TTL - between 1 and 31536000 seconds. + between 1 and 31536000 seconds. `session_affinity_mode` is ``hard`` + (a bound session stays on its worker while that worker is eligible) or + ``soft`` (the binding is a preference the selection policy may override). """ ... diff --git a/lib/kv-router/src/services/selection/affinity/README.md b/lib/kv-router/src/services/selection/affinity/README.md index 8014d7fb373f..dc26bfcd51bc 100644 --- a/lib/kv-router/src/services/selection/affinity/README.md +++ b/lib/kv-router/src/services/selection/affinity/README.md @@ -47,6 +47,10 @@ One writer-id rule: the id installed with the table's replication sink. Frontend Two transports stay: the selection service's ZMQ peer mesh (`services/common/replica_sync.rs`, topics `dynamo.session-affinity.v1` and `.v2`) and the frontend's runtime event plane (`lib/llm/src/session_affinity/replica_sync.rs`, subjects `session_affinity_events` for the old partition-less payload and `session_affinity_events_v2` for the shared schema). For one release both transports publish and apply both versions, so mixed-version replicas converge; applying the same binding twice is idempotent (the second apply refreshes the same version). The v1 forms carry a removal TODO. +## Configuration + +`SessionAffinityConfig` (TTL, mode, entry limit, session-id limit) is the one configuration, and `SessionAffinityConfig::validate` the one validator; `ttl_from_secs_f64` is the one range check for a TTL given in seconds. The selection service builder, the frontend hosts, the Python bindings (`SelectionService`, `KvRouter`, and the `validate_session_affinity_ttl_secs` function the frontend's argument parser calls), and the EPP all go through it. The mode (`hard` or `soft`, `hard` by default) is settable on every host: `--router-session-affinity-mode` on the frontend, `--session-affinity-mode` and `SelectionService(session_affinity_mode=...)` on the standalone service, `DYN_EPP_SESSION_AFFINITY_MODE` on the EPP. + ## Errors Hosts map `AffinityError` onto their own error types. The core maps `InvalidArgument` (the request contradicts the binding or exceeds the session-id limit) to `BadRequest`, `ResourceExhausted` to `NotReady`, `Cancelled` to scheduler shutdown, and `Dropped` to `Internal`. A full table never fails `resolve`: the request routes without affinity and `full_table_fallbacks` counts it (a plain counter on the resolver; not exported as a metric). diff --git a/lib/kv-router/src/services/selection/affinity/table.rs b/lib/kv-router/src/services/selection/affinity/table.rs index 96eb63bf2ebb..a64771df97e8 100644 --- a/lib/kv-router/src/services/selection/affinity/table.rs +++ b/lib/kv-router/src/services/selection/affinity/table.rs @@ -30,6 +30,7 @@ use tokio_util::sync::CancellationToken; pub use crate::protocols::WorkerAffinityTarget as AffinityTarget; +pub const MIN_SESSION_AFFINITY_TTL_SECS: u64 = 1; pub const MAX_SESSION_AFFINITY_TTL_SECS: u64 = 31_536_000; pub const MAX_SESSION_AFFINITY_ENTRIES: usize = 65_536; pub const MAX_SESSION_AFFINITY_ID_BYTES: usize = 256; @@ -121,7 +122,10 @@ enum AffinityEntry { } /// How a table binds sessions; see [`SessionAffinityMode`] for the -/// dispatch-time semantics of `mode`. +/// dispatch-time semantics of `mode`. This is the one session-affinity +/// configuration every host builds, and [`Self::validate`] is its one +/// validator: the standalone service, the frontend, the Python bindings, and +/// the EPP all check a TTL through it. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct SessionAffinityConfig { pub ttl: Duration, @@ -144,6 +148,46 @@ impl SessionAffinityConfig { self.mode = mode; self } + + /// Parse an idle TTL given in seconds, fractional seconds allowed, within + /// the supported range. Non-finite values are rejected. + pub fn ttl_from_secs_f64(secs: f64) -> Result { + if !secs.is_finite() + || secs < MIN_SESSION_AFFINITY_TTL_SECS as f64 + || secs > MAX_SESSION_AFFINITY_TTL_SECS as f64 + { + return Err(invalid_ttl()); + } + Ok(Duration::from_secs_f64(secs)) + } + + /// Check the TTL range and the limits; `with_config` runs this before + /// starting a table. + pub fn validate(&self) -> Result<(), AffinityError> { + if !(Duration::from_secs(MIN_SESSION_AFFINITY_TTL_SECS) + ..=Duration::from_secs(MAX_SESSION_AFFINITY_TTL_SECS)) + .contains(&self.ttl) + { + return Err(invalid_ttl()); + } + if self.max_entries == 0 { + return Err(AffinityError::InvalidArgument( + "session affinity entry limit must be greater than zero".to_string(), + )); + } + if self.max_session_id_bytes == 0 { + return Err(AffinityError::InvalidArgument( + "session affinity ID limit must be greater than zero".to_string(), + )); + } + Ok(()) + } +} + +fn invalid_ttl() -> AffinityError { + AffinityError::InvalidArgument(format!( + "session affinity TTL must be between {MIN_SESSION_AFFINITY_TTL_SECS} and {MAX_SESSION_AFFINITY_TTL_SECS} seconds" + )) } struct Inner { @@ -240,19 +284,8 @@ pub enum AcquireStep { } impl SessionAffinity { - pub(crate) fn validate_ttl(ttl: Duration) -> Result<(), AffinityError> { - if !(Duration::from_secs(1)..=Duration::from_secs(MAX_SESSION_AFFINITY_TTL_SECS)) - .contains(&ttl) - { - return Err(AffinityError::InvalidArgument(format!( - "session affinity TTL must be between 1 and {MAX_SESSION_AFFINITY_TTL_SECS} seconds" - ))); - } - Ok(()) - } - pub fn with_config(config: SessionAffinityConfig) -> Result { - Self::validate_ttl(config.ttl)?; + config.validate()?; let inner = Arc::new(Inner { entries: DashMap::new(), ttl: config.ttl, @@ -1200,6 +1233,44 @@ mod tests { assert_eq!(table.lease_count("s"), Some(0)); } + #[test] + fn config_validates_the_ttl_range_once_for_every_host() { + for secs in [-1.0, 0.0, 0.5, f64::NAN, f64::INFINITY, 31_536_001.0, 1e300] { + let error = SessionAffinityConfig::ttl_from_secs_f64(secs) + .expect_err("out-of-range or non-finite TTLs are rejected"); + assert!(matches!(error, AffinityError::InvalidArgument(_))); + assert!( + error.to_string().contains("session affinity TTL"), + "{error}" + ); + } + for secs in [1.0, 1.5, 31_536_000.0] { + let ttl = SessionAffinityConfig::ttl_from_secs_f64(secs).expect("in range"); + SessionAffinityConfig::new(ttl).validate().expect("valid"); + } + assert!( + SessionAffinityConfig::new(Duration::ZERO) + .validate() + .is_err() + ); + assert!( + SessionAffinityConfig { + max_entries: 0, + ..SessionAffinityConfig::new(TTL) + } + .validate() + .is_err() + ); + assert!( + SessionAffinityConfig { + max_session_id_bytes: 0, + ..SessionAffinityConfig::new(TTL) + } + .validate() + .is_err() + ); + } + #[tokio::test] async fn initializer_cannot_change_the_joined_rank_in_hard_mode() { let table = table(); diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 36367f98c308..6546f67d00fa 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -304,6 +304,8 @@ pub struct SelectionServiceConfig { pub selection_cache: SelectionCacheConfig, /// Session stickiness TTL; `None` disables session affinity. pub session_affinity_ttl: Option, + /// How a bound session treats a dispatch that landed elsewhere. + pub session_affinity_mode: super::affinity::SessionAffinityMode, } type SelectionEntries = RwLock>>>>; diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 057b328dcbf5..af796b3678c1 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -16,7 +16,7 @@ use crate::services::common::replica_sync::{ use crate::services::indexer::backend::IndexerPolicy; use crate::tracking_hash::TrackingHashContext; -use super::affinity::SessionAffinityConfig; +use super::affinity::{SessionAffinityConfig, SessionAffinityMode}; use super::core::{KvIndexSource, SelectionCore, SelectionHost, SelectionServiceConfig}; use super::error::SelectionError; use super::pending::SelectionCacheConfig; @@ -40,6 +40,7 @@ pub struct SelectionServiceBuilder { host: SelectionHost, worker_selection_policy_factory: Option, session_affinity_ttl: Option, + session_affinity_mode: SessionAffinityMode, host_manages_request_lifecycle: bool, } @@ -81,6 +82,7 @@ impl SelectionServiceBuilder { host: SelectionHost::default(), worker_selection_policy_factory: None, session_affinity_ttl: None, + session_affinity_mode: SessionAffinityMode::default(), host_manages_request_lifecycle: false, } } @@ -125,6 +127,13 @@ impl SelectionServiceBuilder { self } + /// How a bound session treats a dispatch that landed elsewhere; `Hard` + /// unless set. Takes effect only with [`Self::session_affinity`]. + pub fn session_affinity_mode(mut self, mode: SessionAffinityMode) -> Self { + self.session_affinity_mode = mode; + self + } + pub fn replica_sync(mut self, port: u16, peers: Vec) -> Self { self.replica_sync_port = Some(port); self.replica_sync_peers = peers; @@ -149,8 +158,11 @@ impl SelectionServiceBuilder { self.kv_router_config .apply_policy_config() .map_err(anyhow::Error::msg)?; - if let Some(ttl) = self.session_affinity_ttl { - super::affinity::SessionAffinity::validate_ttl(ttl)?; + let session_affinity = self + .session_affinity_ttl + .map(|ttl| SessionAffinityConfig::new(ttl).with_mode(self.session_affinity_mode)); + if let Some(config) = &session_affinity { + config.validate()?; } self.kv_router_config .validate_config() @@ -197,7 +209,7 @@ impl SelectionServiceBuilder { self.selection_cache, tracking_hash, indexer_policy, - self.session_affinity_ttl.map(SessionAffinityConfig::new), + session_affinity, )); if recover_from_peers { @@ -267,7 +279,9 @@ impl SelectionServiceConfig { builder = builder.replica_sync(port, self.replica_sync_peers.clone()); } if let Some(ttl) = self.session_affinity_ttl { - builder = builder.session_affinity(ttl); + builder = builder + .session_affinity(ttl) + .session_affinity_mode(self.session_affinity_mode); } builder } diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index ce1d625918e2..2d6cf283db88 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -1778,6 +1778,7 @@ async fn selector_replica_sync_propagates_request_lifecycle() { threads: 1, indexer_peers: Vec::new(), session_affinity_ttl: None, + session_affinity_mode: Default::default(), replica_sync_port: Some(port_a), replica_sync_peers: Vec::new(), kv_router_config: test_config(), diff --git a/lib/llm/src/session_affinity/mod.rs b/lib/llm/src/session_affinity/mod.rs index 77b2ce18850a..cb5eb904f0ab 100644 --- a/lib/llm/src/session_affinity/mod.rs +++ b/lib/llm/src/session_affinity/mod.rs @@ -12,7 +12,8 @@ mod replica_sync; pub(crate) use dynamo_kv_router::services::selection::affinity::Hold; pub use dynamo_kv_router::services::selection::affinity::{ - MAX_SESSION_AFFINITY_TTL_SECS, SessionAffinityMode, + MAX_SESSION_AFFINITY_TTL_SECS, MIN_SESSION_AFFINITY_TTL_SECS, SessionAffinityConfig, + SessionAffinityMode, }; pub use host::{AffinityTarget, HostAffinity, explicit_target}; pub(crate) use host::{