fix(router): unify route_request to async with optional request_json

Signed-off-by: Nejc Habjan <nejc.habjan@siemens.com>

Author: nejch

src/tests/test_roundrobin_router.py

@@ -1,6 +1,8 @@
 import random
 from typing import Dict, List, Tuple
 
+import pytest
+
 from vllm_router.routers.routing_logic import RoundRobinRouter
 
 
@@ -47,21 +49,24 @@ def generate_request(request_type="http") -> Request:
     return Request({"type": request_type})
 
 
-def test_roundrobin_logic(
+@pytest.mark.asyncio
+async def test_roundrobin_logic(
     dynamic_discoveries: int = 10, max_endpoints: int = 1000, max_requests: int = 10000
 ):
     """
     Ensure that all active urls have roughly same number of requests (difference at most 1)
     """
     router = RoundRobinRouter()
 
-    def _fixed_router_check(num_endpoints: int, num_requests: int) -> bool:
+    async def _fixed_router_check(num_endpoints: int, num_requests: int) -> bool:
         # Make num_requests requests to the router and check even output distribution
         endpoints, engine_stats, request_stats = generate_request_args(num_endpoints)
         output_distribution = {}
         for request_idx in range(num_requests):
             request = generate_request()
-            url = router.route_request(endpoints, engine_stats, request_stats, request)
+            url = await router.route_request(
+                endpoints, engine_stats, request_stats, request
+            )
             output_distribution[url] = output_distribution.get(url, 0) + 1
         request_counts = output_distribution.values()
         return max(request_counts) - min(request_counts) <= 1
@@ -70,5 +75,5 @@ def _fixed_router_check(num_endpoints: int, num_requests: int) -> bool:
         num_endpoints = random.randint(1, max_endpoints)
         num_requests = random.randint(1, max_requests)
         # Perform router check
-        res = _fixed_router_check(num_endpoints, num_requests)
+        res = await _fixed_router_check(num_endpoints, num_requests)
         assert res

src/vllm_router/routers/routing_logic.py

@@ -113,15 +113,16 @@ def extract_session_id(self, request: Request, request_json: Dict) -> Optional[s
         if session_key is None:
             return None
         val = request.headers.get(session_key)
-        return val if val is not None else request_json.get(session_key, None)
+        return val if val is not None else (request_json or {}).get(session_key, None)
 
     @abc.abstractmethod
-    def route_request(
+    async def route_request(
         self,
         endpoints: List[EndpointInfo],
         engine_stats: Dict[str, EngineStats],
         request_stats: Dict[str, RequestStats],
         request: Request,
+        request_json: Optional[Dict] = None,
     ) -> str:
         """
         Route the request to the appropriate engine URL
@@ -149,12 +150,13 @@ def __init__(self):
         self.last_endpoints_hash = None
         self._initialized = True
 
-    def route_request(
+    async def route_request(
         self,
         endpoints: List[EndpointInfo],
         engine_stats: Dict[str, EngineStats],
         request_stats: Dict[str, RequestStats],
         request: Request,
+        request_json: Optional[Dict] = None,
     ) -> str:
         """
         Route the request to the appropriate engine URL using a simple
@@ -201,7 +203,7 @@ async def route_request(
         engine_stats: Dict[str, EngineStats],
         request_stats: Dict[str, RequestStats],
         request: Request,
-        request_json: Dict,
+        request_json: Optional[Dict] = None,
     ) -> str:
         """
         Route the request to the appropriate engine URL by the 'session id' in
@@ -320,7 +322,7 @@ async def route_request(
         engine_stats: Dict[str, EngineStats],
         request_stats: Dict[str, RequestStats],
         request: Request,
-        request_json: Dict,
+        request_json: Optional[Dict] = None,
     ) -> str:
         """
         Route the request to the appropriate engine URL by where the KV cache
@@ -346,14 +348,14 @@ async def route_request(
                 self.tokenizer = AutoTokenizer.from_pretrained(
                     endpoints[0].model_names[0]
                 )
-            token_ids = self.tokenizer.encode(request_json.get("prompt", ""))
+            token_ids = self.tokenizer.encode((request_json or {}).get("prompt", ""))
         except Exception:
             # Remote /tokenize fallback (let errors bubble up to keep behavior simple)
             remote_url = endpoints[0].url + "/tokenize"
             headers = {"Content-Type": "application/json"}
             data = {
                 "model": endpoints[0].model_names[0],
-                "prompt": request_json.get("prompt", ""),
+                "prompt": (request_json or {}).get("prompt", ""),
             }
             body = requests.post(
                 remote_url, headers=headers, json=data, timeout=10
@@ -435,7 +437,7 @@ async def route_request(
         engine_stats: Dict[str, EngineStats],
         request_stats: Dict[str, RequestStats],
         request: Request,
-        request_json: Dict,
+        request_json: Optional[Dict] = None,
     ) -> str:
         """
         Route the request to the appropriate engine URL by where the longest
@@ -455,9 +457,10 @@ async def route_request(
         """
 
         # Handle chat completions
-        if "messages" in request_json:
+        _request_json = request_json or {}
+        if "messages" in _request_json:
             # Get the last message from the messages array
-            messages = request_json["messages"]
+            messages = _request_json["messages"]
             if messages:
                 # Concatenate all message content
                 prompt_parts = []
@@ -478,7 +481,7 @@ async def route_request(
                 prompt = ""
         else:
             # Handle regular completions
-            prompt = request_json["prompt"]
+            prompt = _request_json["prompt"]
 
         available_endpoints = set(endpoint.url for endpoint in endpoints)
         _, matched_endpoint = await self.hashtrie.longest_prefix_match(
@@ -503,20 +506,20 @@ def __init__(self, prefill_model_labels: List[str], decode_model_labels: List[st
         self.decode_model_labels = decode_model_labels
         self.request_cache = {}  # Cache to store prefill results
 
-    def route_request(
+    async def route_request(
         self,
         endpoints: List[EndpointInfo],
         engine_stats: Dict[str, EngineStats],
         request_stats: Dict[str, RequestStats],
         request: Request,
-        request_json: Dict,
+        request_json: Optional[Dict] = None,
     ) -> str:
         """
         Route the request to appropriate endpoints for prefill and decode operations.
         First request goes to prefill endpoint, then second request goes to decode endpoint.
         """
         # Find prefill and decode endpoints
-        is_prefill = request_json.get("max_tokens", 0) == 1
+        is_prefill = (request_json or {}).get("max_tokens", 0) == 1
         if is_prefill:
             logger.info("Prefill request")
         else:
@@ -631,7 +634,7 @@ async def route_request(
         engine_stats: Dict[str, EngineStats],
         request_stats: Dict[str, RequestStats],
         request: Request,
-        request_json: Dict,
+        request_json: Optional[Dict] = None,
     ) -> str:
         """
         This method is called by the router framework but for orchestrated routing,

src/vllm_router/services/request_service/request.py

@@ -30,9 +30,6 @@
 from vllm_router.routers.routing_logic import (
     DisaggregatedPrefillOrchestratedRouter,
     DisaggregatedPrefillRouter,
-    KvawareRouter,
-    PrefixAwareRouter,
-    SessionRouter,
 )
 from vllm_router.service_discovery import get_service_discovery
 from vllm_router.services.request_service.rewriter import (
@@ -525,16 +522,10 @@ async def route_general_request(
             f"Routing request {request_id} to engine with Id: {endpoints[0].Id}"
         )
 
-    elif isinstance(
-        request.app.state.router, (KvawareRouter, PrefixAwareRouter, SessionRouter)
-    ):
+    else:
         server_url = await request.app.state.router.route_request(
             endpoints, engine_stats, request_stats, request, request_json
         )
-    else:
-        server_url = request.app.state.router.route_request(
-            endpoints, engine_stats, request_stats, request
-        )
 
     curr_time = time.time()
     # Extract actual session ID from request headers for logging
@@ -572,17 +563,10 @@ async def route_general_request(
                 break
             if request_endpoint:
                 server_url = remaining[0].url
-            elif isinstance(
-                request.app.state.router,
-                (KvawareRouter, PrefixAwareRouter, SessionRouter),
-            ):
+            else:
                 server_url = await request.app.state.router.route_request(
                     remaining, engine_stats, request_stats, request, request_json
                 )
-            else:
-                server_url = request.app.state.router.route_request(
-                    remaining, engine_stats, request_stats, request
-                )
             logger.info(
                 f"Routing request {request_id} to {server_url} "
                 f"(attempt {attempt + 1}/{max_attempts})"
@@ -1189,7 +1173,7 @@ async def proxy_multipart_request(
     request_stats = request_stats_monitor.get_request_stats(time.time())
 
     # pick one using the router's configured logic (roundrobin, least-loaded, etc.)
-    chosen_url = router.route_request(
+    chosen_url = await router.route_request(
         endpoints,
         engine_stats,
         request_stats,