diff --git a/src/tests/test_request_auth_headers.py b/src/tests/test_request_auth_headers.py index 17cedcecd..786401f52 100644 --- a/src/tests/test_request_auth_headers.py +++ b/src/tests/test_request_auth_headers.py @@ -32,7 +32,7 @@ def _build_request(headers=None): state.callbacks = None state.request_stats_monitor.get_request_stats.return_value = {} state.engine_stats_scraper.get_engine_stats.return_value = {} - state.router.route_request.return_value = "http://whisper-engine" + state.router.route_request = AsyncMock(return_value="http://whisper-engine") request = MagicMock() request.headers = headers or {} diff --git a/src/tests/test_roundrobin_router.py b/src/tests/test_roundrobin_router.py index 3a649fafb..8cbe27c5b 100644 --- a/src/tests/test_roundrobin_router.py +++ b/src/tests/test_roundrobin_router.py @@ -62,7 +62,8 @@ def assert_even_distribution(route_counts: Counter): assert max(counts) - min(counts) <= 1 -def test_roundrobin_logic( +@pytest.mark.asyncio +async def test_roundrobin_logic( dynamic_discoveries: int = 10, max_endpoints: int = 1000, max_requests: int = 10000 ): """ @@ -70,12 +71,14 @@ def test_roundrobin_logic( """ router = RoundRobinRouter() - def assert_router_stays_balanced(num_endpoints: int, num_requests: int): + async def assert_router_stays_balanced(num_endpoints: int, num_requests: int): endpoints, engine_stats, request_stats = generate_request_args(num_endpoints) route_counts = Counter() for _ 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 + ) route_counts[url] += 1 assert_even_distribution(route_counts) @@ -83,10 +86,11 @@ def assert_router_stays_balanced(num_endpoints: int, num_requests: int): for _ in range(dynamic_discoveries): num_endpoints = random.randint(1, max_endpoints) num_requests = random.randint(1, max_requests) - assert_router_stays_balanced(num_endpoints, num_requests) + await assert_router_stays_balanced(num_endpoints, num_requests) -def test_roundrobin_keeps_state_per_endpoint_set(): +@pytest.mark.asyncio +async def test_roundrobin_keeps_state_per_endpoint_set(): router = RoundRobinRouter() request = generate_request() engine_stats = {} @@ -97,16 +101,21 @@ def test_roundrobin_keeps_state_per_endpoint_set(): route_counts_b = Counter() for _ in range(100): - url_a = router.route_request(endpoints_a, engine_stats, request_stats, request) + url_a = await router.route_request( + endpoints_a, engine_stats, request_stats, request + ) route_counts_a[url_a] += 1 - url_b = router.route_request(endpoints_b, engine_stats, request_stats, request) + url_b = await router.route_request( + endpoints_b, engine_stats, request_stats, request + ) route_counts_b[url_b] += 1 assert_even_distribution(route_counts_a) assert_even_distribution(route_counts_b) -def test_roundrobin_keeps_state_when_endpoint_order_changes(): +@pytest.mark.asyncio +async def test_roundrobin_keeps_state_when_endpoint_order_changes(): router = RoundRobinRouter() request = generate_request() engine_stats = {} @@ -115,10 +124,11 @@ def test_roundrobin_keeps_state_when_endpoint_order_changes(): endpoints_a_reordered = [EndpointInfo(url="a1"), EndpointInfo(url="a0")] assert ( - router.route_request(endpoints_a, engine_stats, request_stats, request) == "a0" + await router.route_request(endpoints_a, engine_stats, request_stats, request) + == "a0" ) assert ( - router.route_request( + await router.route_request( endpoints_a_reordered, engine_stats, request_stats, @@ -128,8 +138,9 @@ def test_roundrobin_keeps_state_when_endpoint_order_changes(): ) -def test_roundrobin_rejects_empty_endpoint_list(): +@pytest.mark.asyncio +async def test_roundrobin_rejects_empty_endpoint_list(): router = RoundRobinRouter() with pytest.raises(ValueError, match="at least one endpoint"): - router.route_request([], {}, {}, generate_request()) + await router.route_request([], {}, {}, generate_request()) diff --git a/src/vllm_router/routers/routing_logic.py b/src/vllm_router/routers/routing_logic.py index 1f56bffb5..d483695db 100644 --- a/src/vllm_router/routers/routing_logic.py +++ b/src/vllm_router/routers/routing_logic.py @@ -115,12 +115,13 @@ def extract_session_id(self, request: Request, request_json: Dict) -> Optional[s return val if val is not None else request_json.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: Dict = {}, ) -> str: """ Route the request to the appropriate engine URL @@ -132,6 +133,7 @@ def route_request( request_stats (Dict[str, RequestStats]): The request stats indicating the request-level performance of each engine request (Request): The incoming request + request_json (Dict): The request body """ raise NotImplementedError @@ -165,12 +167,13 @@ def _endpoint_key(self, endpoints: List[EndpointInfo]) -> tuple[str, ...]: self._sorted_cache[urls] = key return key - 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 = {}, ) -> str: """ Route the request to the appropriate engine URL using a simple @@ -183,6 +186,7 @@ def route_request( request_stats (Dict[str, RequestStats]): The request stats indicating the request-level performance of each engine request (Request): The incoming request + request_json (Dict): Unused for RoundRobinRouter """ endpoint_urls = self._endpoint_key(endpoints) idx = self._next_index.get(endpoint_urls, 0) @@ -493,7 +497,7 @@ async def route_request( prompt = "" else: # Handle regular completions - prompt = request_json["prompt"] + prompt = request_json.get("prompt", "") available_endpoints = set(endpoint.url for endpoint in endpoints) _, matched_endpoint = await self.hashtrie.longest_prefix_match( @@ -518,7 +522,7 @@ 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], diff --git a/src/vllm_router/services/request_service/request.py b/src/vllm_router/services/request_service/request.py index 90bb0ec20..623838ab1 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 ( @@ -554,17 +551,10 @@ async def route_general_request( logger.debug( 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 @@ -602,17 +592,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})" @@ -1110,6 +1093,7 @@ async def route_general_transcriptions( ) language: Optional[str] = form.get("language", "en") stream: bool = form.get("stream", "false").lower() == "true" + session_id: Optional[str] = form.get("session_id", None) except KeyError as e: return JSONResponse( status_code=400, @@ -1119,13 +1103,14 @@ async def route_general_transcriptions( logger.debug("==== Enter audio_transcriptions ====") logger.debug("Received upload: %s (%s)", file.filename, file.content_type) logger.debug( - "Params: model=%s prompt=%r response_format=%r temperature=%r language=%s stream=%s", + "Params: model=%s prompt=%r response_format=%r temperature=%r language=%s stream=%s session_id=%s", model, prompt, response_format, temperature, language, stream, + session_id, ) payload_bytes = await file.read() @@ -1145,6 +1130,9 @@ async def route_general_transcriptions( if stream: data["stream"] = "true" + if session_id: + data["session_id"] = str(session_id) + form_data = aiohttp.FormData() for key, (filename, content, content_type) in files.items(): @@ -1154,7 +1142,7 @@ async def route_general_transcriptions( form_data.add_field(key, value) return await proxy_multipart_request( - form_data, model, endpoint, request, stream=stream + form_data, model, endpoint, request, data, stream=stream ) @@ -1185,6 +1173,7 @@ async def proxy_multipart_request( model: str, endpoint: str, request: Request, + request_json: dict = {}, *, stream: bool = False, ): @@ -1218,11 +1207,8 @@ 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( - endpoints, - engine_stats, - request_stats, - request, + chosen_url = await router.route_request( + endpoints, engine_stats, request_stats, request, request_json ) logger.info( "Proxying multi-part form request for model %s to %s", model, chosen_url