[Bugfix] Fix /v1/audio/transcriptions endpoint when not using roundrobin

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 {}

src/tests/test_roundrobin_router.py

@@ -62,31 +62,35 @@ 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
 ):
     """
     Ensure that all active urls have roughly same number of requests (difference at most 1)
     """
     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)
 
     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())

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],

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