feat(kv-router): enforce hard session affinity in the scheduler for every policy

docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx

@@ -379,7 +379,7 @@ Routing cache resources start only if a component requests `CACHE`. `serve_index
 | `block_size()` | Tokens in one KV block |
 | `tracks_prefill_tokens()` | Whether the request contributes to prefill-load tracking |
 | `session_context()` | Optional session metadata described below |
-| `affinity_target()` | Optional advisory worker and data-parallel rank from soft session affinity |
+| `affinity_target()` | Optional session-affinity worker and data-parallel rank. Under hard affinity the host already limits the candidate set to it while it is eligible; under soft affinity it is advisory |
 | `expected_output_tokens()` | Optional expected output length |
 | `priority_jump()` | Scheduler priority boost. Queue policies treat negative values as zero |
 | `strict_priority()` | Strict integer priority. The queue orders larger values first |

examples/router/custom-policy-example/disagg-filter-score-pick/src/lib.rs

@@ -153,7 +153,7 @@ mod tests {
             isl_tokens: 16,
             lora_name: None,
             expected_output_tokens: None,
-            affinity_target: None,
+            affinity: None,
             pinned_worker: None,
             allowed_worker_ids: None,
             routing_constraints: RoutingConstraints::default(),

examples/router/custom-policy-example/soft-pin-repin/src/lib.rs

@@ -123,7 +123,7 @@ mod tests {
     use std::collections::HashMap;
 
     use dynamo_kv_router::protocols::{RoutingConstraints, WorkerConfigLike, WorkerWithDpRank};
-    use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode};
+    use dynamo_kv_router::scheduling::{AffinityRequirement, OverlapSignals, ScheduleMode};
     use dynamo_kv_router::{
         SchedulingRequest, WorkerLoadProjection, WorkerSelectionInput, WorkerSelector,
     };
@@ -157,7 +157,7 @@ mod tests {
             isl_tokens: 16,
             lora_name: None,
             expected_output_tokens: None,
-            affinity_target,
+            affinity: affinity_target.map(AffinityRequirement::soft),
             pinned_worker: None,
             allowed_worker_ids: None,
             routing_constraints: RoutingConstraints::default(),

lib/kv-router/benches/worker_selection.rs

@@ -192,7 +192,7 @@ fn fixture_with_preferred_taints(
         isl_tokens: 2_048,
         lora_name: None,
         expected_output_tokens: Some(256),
-        affinity_target: None,
+        affinity: None,
         pinned_worker: None,
         allowed_worker_ids: None,
         routing_constraints: RoutingConstraints {

lib/kv-router/src/lib.rs

@@ -71,8 +71,8 @@ pub use scheduling::LoraWorkerFilter;
 pub use scheduling::PrefillLoadEstimator;
 pub use scheduling::policy::{FcfsPolicy, RouterSchedulingPolicy, SchedulingPolicy, WsptPolicy};
 pub use scheduling::{
-    KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse, SessionContext,
-    WorkerSelectionInputTrigger, WorkerSelectionPolicyError,
+    AffinityRequirement, AffinityStrength, KvSchedulerError, PotentialLoad, SchedulingRequest,
+    SchedulingResponse, SessionContext, WorkerSelectionInputTrigger, WorkerSelectionPolicyError,
 };
 #[cfg(any(test, feature = "bench"))]
 pub use selector::DefaultWorkerSelector;

lib/kv-router/src/plugins/worker_selection/context.rs

@@ -18,9 +18,10 @@ pub struct WorkerSelectionContext<'a> {
 }
 
 impl WorkerSelectionContext<'_> {
-    /// The exact worker/rank imposed by the host for this selection, if any.
-    /// Includes explicit pins and eligible exclusive-affinity targets. This is
-    /// read-only routing metadata, not permission to change eligibility.
+    /// The exact worker/rank the request pinned explicitly, if any. This is
+    /// read-only routing metadata, not permission to change eligibility. A
+    /// `Hard` session binding is not reported here; the host limits the
+    /// candidate set to it instead.
     pub fn pinned_worker(&self) -> Option<WorkerWithDpRank> {
         self.pinned_worker
     }
@@ -55,10 +56,12 @@ impl WorkerSelectionContext<'_> {
 
     /// Return the session-affinity target resolved by the request host.
     ///
-    /// The default selector treats an eligible target as exclusive. Custom policies receive it as
-    /// advisory context; it may be absent from their candidate set when unavailable or filtered.
+    /// Under `Hard` session affinity the host limits every policy's candidates to this target
+    /// while it is eligible, so the candidate set already reflects it. Under `Soft` affinity it is
+    /// a preference: the policy receives the full eligible candidate set and may select another
+    /// worker. The target may be absent from the candidate set when it is unavailable or filtered.
     pub fn affinity_target(&self) -> Option<WorkerAffinityTarget> {
-        self.request.affinity_target
+        self.request.affinity.map(|affinity| affinity.target)
     }
 
     /// Return the expected output length, if the request supplies one.

lib/kv-router/src/protocols.rs

@@ -464,6 +464,11 @@ impl WorkerAffinityTarget {
     pub fn new(worker_id: WorkerId, dp_rank: Option<DpRank>) -> Self {
         Self { worker_id, dp_rank }
     }
+
+    /// Whether `worker` is this target: the worker, and the rank when bound.
+    pub fn matches(self, worker: WorkerWithDpRank) -> bool {
+        self.worker_id == worker.worker_id && self.dp_rank.is_none_or(|rank| rank == worker.dp_rank)
+    }
 }
 
 impl From<WorkerWithDpRank> for WorkerAffinityTarget {

lib/kv-router/src/scheduling/filter.rs

@@ -186,11 +186,9 @@ impl<'a> RoutingEligibility<'a> {
                 worker_id: worker.worker_id,
             });
         }
-        if !self.matches_affinity_target(worker.worker_id)
-            || self
-                .affinity_target
-                .and_then(|target| target.dp_rank)
-                .is_some_and(|rank| rank != worker.dp_rank)
+        if !self
+            .affinity_target
+            .is_none_or(|target| target.matches(worker))
         {
             return Err(WorkerEligibilityError::WorkerNotAllowed {
                 worker_id: worker.worker_id,
@@ -670,6 +668,56 @@ mod tests {
         );
     }
 
+    #[test]
+    fn routing_eligibility_ranked_affinity_target_yields_exact_rank() {
+        let workers = workers();
+        let constraints = RoutingConstraints::default();
+        let eligibility = RoutingEligibility::new(None, None, None, &constraints)
+            .with_affinity_target(WorkerAffinityTarget::new(7, Some(3)));
+        let mut ranks = Vec::new();
+
+        eligibility.for_each_eligible_worker_rank(&workers, |worker, _| ranks.push(worker));
+
+        assert_eq!(ranks, vec![WorkerWithDpRank::new(7, 3)]);
+    }
+
+    #[test]
+    fn affinity_target_is_eligible_requires_a_schedulable_target() {
+        let workers = workers();
+        let constraints = RoutingConstraints::default();
+        let open = RoutingEligibility::new(None, None, None, &constraints);
+        let target = WorkerAffinityTarget::new(7, None);
+
+        assert!(open.affinity_target_is_eligible(&workers, target));
+        assert!(open.affinity_target_is_eligible(&workers, WorkerAffinityTarget::new(7, Some(4))));
+        assert!(
+            !open.affinity_target_is_eligible(&workers, WorkerAffinityTarget::new(7, Some(5))),
+            "rank out of range"
+        );
+        assert!(
+            !open.affinity_target_is_eligible(&workers, WorkerAffinityTarget::new(8, None)),
+            "unknown worker"
+        );
+
+        let others = HashSet::from([8]);
+        let disallowed = RoutingEligibility::new(Some(&others), None, None, &constraints);
+        assert!(
+            !disallowed.affinity_target_is_eligible(&workers, target),
+            "outside the caller's set"
+        );
+        let overloaded = HashSet::from([7]);
+        let shed = RoutingEligibility::new(None, Some(&overloaded), None, &constraints);
+        assert!(
+            !shed.affinity_target_is_eligible(&workers, target),
+            "overloaded"
+        );
+        let unavailable = open.with_available_workers(Some(&others));
+        assert!(
+            !unavailable.affinity_target_is_eligible(&workers, target),
+            "unavailable"
+        );
+    }
+
     #[test]
     fn routing_eligibility_pinned_expansion_yields_exact_rank() {
         let workers = workers();

lib/kv-router/src/scheduling/local.rs

@@ -88,7 +88,7 @@ where
             isl_tokens,
             lora_name,
             expected_output_tokens,
-            affinity_target,
+            affinity,
             pinned_worker,
             allowed_worker_ids,
             routing_constraints,
@@ -108,7 +108,7 @@ where
             isl_tokens,
             lora_name,
             expected_output_tokens,
-            affinity_target,
+            affinity,
             pinned_worker,
             allowed_worker_ids,
             routing_constraints,
@@ -582,7 +582,7 @@ where
             policy_class,
             session_context: None,
             expected_output_tokens,
-            affinity_target: None,
+            affinity: None,
             pinned_worker,
             allowed_worker_ids,
             shared_cache_hits,
@@ -1085,7 +1085,7 @@ mod tests {
             isl_tokens: 64,
             lora_name: None,
             expected_output_tokens: None,
-            affinity_target: None,
+            affinity: None,
             pinned_worker: None,
             allowed_worker_ids: None,
             routing_constraints: crate::protocols::RoutingConstraints::default(),

lib/kv-router/src/scheduling/policy.rs

@@ -223,7 +223,7 @@ mod tests {
             policy_class: None,
             session_context: None,
             expected_output_tokens: None,
-            affinity_target: None,
+            affinity: None,
             pinned_worker: None,
             allowed_worker_ids: None,
             routing_constraints: crate::protocols::RoutingConstraints::default(),