From 4b09be866632e82ded759dc4513a34c093408405 Mon Sep 17 00:00:00 2001 From: Nejc Habjan Date: Thu, 23 Apr 2026 08:50:27 +0200 Subject: [PATCH] fix(router): unify route_request to async with optional request_json Signed-off-by: Nejc Habjan --- src/tests/test_roundrobin_router.py | 13 +++++--- src/vllm_router/routers/routing_logic.py | 33 ++++++++++--------- .../services/request_service/request.py | 22 ++----------- 3 files changed, 30 insertions(+), 38 deletions(-) diff --git a/src/tests/test_roundrobin_router.py b/src/tests/test_roundrobin_router.py index d336cef55..39a1d581b 100644 --- a/src/tests/test_roundrobin_router.py +++ b/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,7 +49,8 @@ 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 ): """ @@ -55,13 +58,15 @@ def test_roundrobin_logic( """ 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 diff --git a/src/vllm_router/routers/routing_logic.py b/src/vllm_router/routers/routing_logic.py index 5659b987a..9bc9da239 100644 --- a/src/vllm_router/routers/routing_logic.py +++ b/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, diff --git a/src/vllm_router/services/request_service/request.py b/src/vllm_router/services/request_service/request.py index 961bf8c15..82280f8d8 100644 --- a/src/vllm_router/services/request_service/request.py +++ b/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,