feat: add sticky least-loaded session routing

Dockerfile.router

@@ -12,7 +12,10 @@ RUN cargo build --release --bin vllm-router
 # Stage 2: Minimal runtime image
 FROM debian:bookworm-slim
 
-RUN apt-get update && apt-get install -y ca-certificates && rm -rf /var/lib/apt/lists/*
+RUN apt-get update \
+    && apt-get upgrade -y \
+    && apt-get install -y --no-install-recommends ca-certificates \
+    && rm -rf /var/lib/apt/lists/*
 
 COPY --from=builder /build/target/release/vllm-router /usr/local/bin/vllm-router
 

README.md

@@ -206,6 +206,7 @@ The router supports multiple load balancing policies:
 | `round_robin` | Sequential distribution across workers | No | General purpose, even distribution |
 | `random` | Uniform random selection | No | Simple deployments |
 | `consistent_hash` | Routes same session/user to same worker | Yes | Multi-turn chat, KV cache reuse |
+| `sticky_least_loaded` | Assigns new sessions to the worker with the fewest active sessions | Yes | Multi-turn workloads with explicit session lifetimes |
 | `power_of_two` | Picks least loaded of two random workers | No | Load-sensitive workloads |
 | `cache_aware` | Optimizes for prefix cache hits | Yes | Repeated prompts, few-shot |
 
@@ -219,6 +220,50 @@ curl -X POST http://router:8000/v1/chat/completions \
 
 For detailed configuration options, hash key priorities, and usage examples, see [Load Balancing Documentation](docs/load_balancing/README.md).
 
+#### Sticky least-loaded sessions
+
+Use `--policy sticky_least_loaded` to balance **active sessions**, not concurrent
+requests, token counts, or GPU utilization. New sessions reserve a worker
+atomically; ties use rendezvous hashing. Existing sessions keep their healthy
+worker, even if session counts later become uneven. Unavailable workers cause
+reassignment on the next request. Prefill and decode pools track sessions separately.
+
+Session IDs prefer stable headers (`x-session-id`, `x-user-id`, `x-tenant-id`,
+and `x-correlation-id`), then the JSON fields `session_params.session_id`, `user`,
+`session_id`, and `user_id`. Per-request `x-request-id` and `x-trace-id` headers
+remain available to `consistent_hash` but do not reserve sticky sessions. Requests
+without an explicit session identifier do not reserve a session.
+
+Release a session after its final request completes:
+
+```bash
+curl -X POST 'http://router:8000/finish_session?session_id=my-session-123'
+```
+
+This endpoint uses the router's configured API-key validation. It releases the ID
+across default, per-model, prefill, and decode policies. Repeated or unknown IDs
+are harmless. It does not cancel generation; use globally unique IDs and finish
+only after all requests for that session have completed. Models using the shared
+default policy must include the model in their session ID (for example,
+`model-a:conversation-123`), or use separate per-model policy instances. Reusing
+one ID across disjoint model pools reassigns its single reservation between pools.
+
+Idle sessions expire after two hours. Set
+`VLLM_ROUTER_SLL_SESSION_EXPIRATION_IN_S` on the router to override the TTL; expired
+entries are swept on routing requests at most once per minute. Revisited sessions
+are checked for expiry before refreshing their affinity. State is local to one
+router process and is lost on restart. Multiple routers need consistent ingress
+routing and session release on each router that tracked the session.
+The policy keeps one entry per active session and serializes assignment under a
+mutex, so callers should release sessions promptly and choose a TTL appropriate
+for their workload.
+
+Each policy instance retains at most 100,000 sessions and accepts session IDs up
+to 256 bytes. Override these limits with `VLLM_ROUTER_SLL_MAX_SESSIONS` and
+`VLLM_ROUTER_SLL_MAX_SESSION_ID_BYTES`. Requests that introduce a session beyond
+either limit continue to route least-loaded without affinity; existing tracked
+sessions remain sticky.
+
 ## Advanced Features
 
 ### Kubernetes Service Discovery

benches/request_processing.rs

@@ -20,6 +20,9 @@ fn default_generate_request() -> GenerateRequest {
         // VLLM Extensions
         lora_path: None,
         session_params: None,
+        user: None,
+        session_id: None,
+        user_id: None,
         return_hidden_states: false,
         rid: None,
     }

py_src/vllm_router/router.py

@@ -15,6 +15,7 @@ def policy_from_str(policy_str: Optional[str]) -> PolicyType:
         "cache_aware": PolicyType.CacheAware,
         "power_of_two": PolicyType.PowerOfTwo,
         "consistent_hash": PolicyType.ConsistentHash,
+        "sticky_least_loaded": PolicyType.StickyLeastLoaded,
     }
     return policy_map[policy_str]
 

py_src/vllm_router/router_args.py

@@ -140,6 +140,7 @@ def add_cli_args(
                 "cache_aware",
                 "power_of_two",
                 "consistent_hash",
+                "sticky_least_loaded",
             ],
             help="Load balancing policy to use. In PD mode, this is used for both prefill and decode unless overridden",
         )
@@ -153,6 +154,7 @@ def add_cli_args(
                 "cache_aware",
                 "power_of_two",
                 "consistent_hash",
+                "sticky_least_loaded",
             ],
             help="Specific policy for prefill nodes in PD mode. If not specified, uses the main policy",
         )
@@ -166,6 +168,7 @@ def add_cli_args(
                 "cache_aware",
                 "power_of_two",
                 "consistent_hash",
+                "sticky_least_loaded",
             ],
             help="Specific policy for decode nodes in PD mode. If not specified, uses the main policy",
         )

py_test/unit/test_arg_parser.py

@@ -409,6 +409,7 @@ def test_valid_policies(self):
         assert policy_from_str("cache_aware") == PolicyType.CacheAware
         assert policy_from_str("power_of_two") == PolicyType.PowerOfTwo
         assert policy_from_str("consistent_hash") == PolicyType.ConsistentHash
+        assert policy_from_str("sticky_least_loaded") == PolicyType.StickyLeastLoaded
 
     def test_invalid_policy(self):
         """Test conversion of invalid policy string."""
@@ -419,6 +420,21 @@ def test_invalid_policy(self):
 class TestParseRouterArgs:
     """Test the parse_router_args function."""
 
+    def test_parse_sticky_least_loaded_for_each_pool(self):
+        args = parse_router_args(
+            [
+                "--policy",
+                "sticky_least_loaded",
+                "--prefill-policy",
+                "sticky_least_loaded",
+                "--decode-policy",
+                "sticky_least_loaded",
+            ]
+        )
+        assert args.policy == "sticky_least_loaded"
+        assert args.prefill_policy == "sticky_least_loaded"
+        assert args.decode_policy == "sticky_least_loaded"
+
     def test_parse_basic_args(self):
         """Test parsing basic router arguments."""
         args = [

src/config/types.rs

@@ -247,6 +247,14 @@ pub enum PolicyConfig {
         virtual_nodes: u32,
     },
 
+    /// Session-aware routing that combines session stickiness with load
+    /// balancing: existing sessions stick to their assigned replica, new
+    /// sessions are routed to the least-loaded replica (ties broken by
+    /// consistent hashing). Session expiry is configured via the
+    /// `VLLM_ROUTER_SLL_SESSION_EXPIRATION_IN_S` environment variable.
+    #[serde(rename = "sticky_least_loaded")]
+    StickyLeastLoaded,
+
     #[serde(rename = "rendezvous_hash")]
     RendezvousHash,
 }
@@ -259,6 +267,7 @@ impl PolicyConfig {
             PolicyConfig::CacheAware { .. } => "cache_aware",
             PolicyConfig::PowerOfTwo { .. } => "power_of_two",
             PolicyConfig::ConsistentHash { .. } => "consistent_hash",
+            PolicyConfig::StickyLeastLoaded => "sticky_least_loaded",
             PolicyConfig::RendezvousHash => "rendezvous_hash",
         }
     }

src/config/validation.rs

@@ -188,6 +188,10 @@ impl ConfigValidator {
                     });
                 }
             }
+            PolicyConfig::StickyLeastLoaded => {
+                // Session expiration is sourced from the environment; no
+                // structured config to validate here.
+            }
             PolicyConfig::RendezvousHash => {
                 // No specific validation needed
             }

src/lib.rs

@@ -27,6 +27,7 @@ pub enum PolicyType {
     CacheAware,
     PowerOfTwo,
     ConsistentHash,
+    StickyLeastLoaded,
 }
 
 #[pyclass]
@@ -125,6 +126,7 @@ impl Router {
                 PolicyType::ConsistentHash => ConfigPolicyConfig::ConsistentHash {
                     virtual_nodes: 160, // Default value
                 },
+                PolicyType::StickyLeastLoaded => ConfigPolicyConfig::StickyLeastLoaded,
             }
         };
 

src/main.rs

@@ -107,7 +107,7 @@ struct CliArgs {
     worker_urls: Vec<String>,
 
     /// Load balancing policy to use
-    #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "rendezvous_hash"])]
+    #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "sticky_least_loaded", "rendezvous_hash"])]
     policy: String,
 
     /// Enable vLLM PD (Prefill-Decode) disaggregated mode with vLLM-specific two-stage processing
@@ -124,11 +124,11 @@ struct CliArgs {
     decode: Vec<String>,
 
     /// Specific policy for prefill nodes in PD mode
-    #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "rendezvous_hash"])]
+    #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "sticky_least_loaded", "rendezvous_hash"])]
     prefill_policy: Option<String>,
 
     /// Specific policy for decode nodes in PD mode
-    #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "rendezvous_hash"])]
+    #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "sticky_least_loaded", "rendezvous_hash"])]
     decode_policy: Option<String>,
 
     /// Timeout in seconds for worker startup
@@ -378,6 +378,7 @@ impl CliArgs {
             "consistent_hash" => PolicyConfig::ConsistentHash {
                 virtual_nodes: 160, // Default value
             },
+            "sticky_least_loaded" => PolicyConfig::StickyLeastLoaded,
             "rendezvous_hash" => PolicyConfig::RendezvousHash,
             _ => PolicyConfig::RoundRobin, // Fallback
         }

src/policies/factory.rs

@@ -3,6 +3,7 @@
 use super::{
     CacheAwareConfig, CacheAwarePolicy, ConsistentHashPolicy, LoadBalancingPolicy,
     PowerOfTwoPolicy, RandomPolicy, RendezvousHashPolicy, RoundRobinPolicy,
+    StickyLeastLoadedPolicy,
 };
 use crate::config::PolicyConfig;
 use std::sync::Arc;
@@ -38,6 +39,7 @@ impl PolicyFactory {
                 // The consistent hash policy uses a hardcoded value for now
                 Arc::new(ConsistentHashPolicy::new())
             }
+            PolicyConfig::StickyLeastLoaded => Arc::new(StickyLeastLoadedPolicy::new()),
             PolicyConfig::RendezvousHash => Arc::new(RendezvousHashPolicy::new()),
         }
     }
@@ -50,6 +52,9 @@ impl PolicyFactory {
             "power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())),
             "cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())),
             "consistent_hash" | "consistenthash" => Some(Arc::new(ConsistentHashPolicy::new())),
+            "sticky_least_loaded" | "stickyleastloaded" => {
+                Some(Arc::new(StickyLeastLoadedPolicy::new()))
+            }
             "rendezvous_hash" | "rendezvoushash" => Some(Arc::new(RendezvousHashPolicy::new())),
             _ => None,
         }
@@ -94,6 +99,10 @@ mod tests {
         // Test RendezvousHash
         let policy = PolicyFactory::create_from_config(&PolicyConfig::RendezvousHash);
         assert_eq!(policy.name(), "rendezvous_hash");
+
+        // Test StickyLeastLoaded
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::StickyLeastLoaded);
+        assert_eq!(policy.name(), "sticky_least_loaded");
     }
 
     #[test]
@@ -108,6 +117,8 @@ mod tests {
         assert!(PolicyFactory::create_by_name("CacheAware").is_some());
         assert!(PolicyFactory::create_by_name("consistent_hash").is_some());
         assert!(PolicyFactory::create_by_name("ConsistentHash").is_some());
+        assert!(PolicyFactory::create_by_name("sticky_least_loaded").is_some());
+        assert!(PolicyFactory::create_by_name("StickyLeastLoaded").is_some());
         assert!(PolicyFactory::create_by_name("rendezvous_hash").is_some());
         assert!(PolicyFactory::create_by_name("RendezvousHash").is_some());
         assert!(PolicyFactory::create_by_name("unknown").is_none());