[Router] Add round robin for session routing

src/vllm_router/app.py

@@ -131,7 +131,11 @@ def initialize_all(app: FastAPI, args):
             args.batch_processor, args.file_storage_path, app.state.batch_storage
         )
 
-    initialize_routing_logic(args.routing_logic, session_key=args.session_key)
+    initialize_routing_logic(
+        args.routing_logic,
+        session_key=args.session_key,
+        session_logic=args.session_logic,
+    )
 
     # Initialize feature gates
     initialize_feature_gates(args.feature_gates)

src/vllm_router/parsers/parser.py

@@ -108,6 +108,13 @@ def parse_args():
         default=None,
         help="The key (in the header) to identify a session.",
     )
+    parser.add_argument(
+        "--session-logic",
+        type=str,
+        default="hashring",
+        choices=["roundrobin", "hashring"],
+        help="The routing logic to apply for the initial requests when using session routing.",
+    )
 
     # Batch API
     # TODO(gaocegege): Make these batch api related arguments to a separate config.

src/vllm_router/routers/routing_logic.py

@@ -82,12 +82,13 @@ class SessionRouter(RoutingInterface):
     in the request headers
     """
 
-    def __init__(self, session_key: str = None):
+    def __init__(self, session_key: str = None, session_logic: str = "hashring"):
         if hasattr(self, "_initialized"):
             return
         if session_key is None:
             raise ValueError("SessionRouter must be initialized with a session_key")
         self.session_key = session_key
+        self.session_logic = session_logic
         self.hash_ring = HashRing()
         self._initialized = True
 
@@ -166,8 +167,15 @@ def route_request(
             # Route based on QPS if no session ID is present
             url = self._qps_routing(endpoints, request_stats)
         else:
-            # Use the hash ring to get the endpoint for the session ID
-            url = self.hash_ring.get_node(session_id)
+            if self.session_logic == "hashring":
+                # Use the hash ring to get the endpoint for the session ID
+                url = self.hash_ring.get_node(session_id)
+            elif self.session_logic == "roundrobin":
+                len_engines = len(endpoints)
+                chosen = sorted(endpoints, key=lambda e: e.url)[
+                    int(session_id) % len_engines
+                ]
+                url = chosen.url
 
         return url
 
@@ -181,7 +189,7 @@ def initialize_routing_logic(
         return RoundRobinRouter()
     elif routing_logic == RoutingLogic.SESSION_BASED:
         logger.info(f"Initializing session-based routing logic with kwargs: {kwargs}")
-        return SessionRouter(kwargs.get("session_key"))
+        return SessionRouter(kwargs.get("session_key"), kwargs.get("session_logic"))
     else:
         raise ValueError(f"Invalid routing logic {routing_logic}")