Adding support to route a request to a specific engine instance

Signed-off-by: Braulio Dumba <Braulio.Dumba@ibm.com>

Author: dumb0002

src/vllm_router/routers/main_router.py

@@ -145,6 +145,39 @@ async def show_models():
     return JSONResponse(content=model_list.model_dump())
 
 
+@main_router.get("/engines")
+async def get_engine_instances():
+    """
+    Returns a list of all models available in the stack.
+
+    Args:
+        None
+
+    Returns:
+        JSONResponse: A JSON response containing the list of models.
+
+    Raises:
+        Exception: If there is an error in retrieving the endpoint information.
+    """
+    endpoints = get_service_discovery().get_endpoint_info()
+    existing_engines = set()
+    engines_cards = []
+
+    for endpoint in endpoints:
+        if endpoint.Id in existing_engines:
+            continue
+        engine_card = {
+            "engine_id": endpoint.Id,
+            "serving_models": endpoint.model_names,
+            "created": endpoint.added_timestamp,
+        }
+
+        engines_cards.append(engine_card)
+        existing_engines.add(endpoint.Id)
+
+    return JSONResponse(content=engines_cards)
+
+
 @main_router.get("/health")
 async def health() -> Response:
     """

src/vllm_router/service_discovery.py

@@ -19,6 +19,7 @@
 import os
 import threading
 import time
+import uuid
 from dataclasses import dataclass
 from typing import Dict, List, Optional
 
@@ -46,6 +47,9 @@ class EndpointInfo:
     # Model names
     model_names: List[str]
 
+    # Endpoint Id
+    Id: str
+
     # Added timestamp
     added_timestamp: float
 
@@ -175,6 +179,7 @@ def __init__(
         self.aliases = aliases
         self.model_labels = model_labels
         self.model_types = model_types
+        self.engines_id = [str(uuid.uuid4()) for i in range(0, len(urls))]
         self.added_timestamp = int(time.time())
         self.unhealthy_endpoint_hashes = []
         if static_backend_health_checks:
@@ -248,6 +253,7 @@ def get_endpoint_info(self) -> List[EndpointInfo]:
             endpoint_info = EndpointInfo(
                 url=url,
                 model_names=[model],  # Convert single model to list
+                Id=self.engines_id[i],
                 added_timestamp=self.added_timestamp,
                 model_label=model_label,
                 model_info=self._get_model_info(model),
@@ -439,6 +445,7 @@ def _add_engine(
                 url=f"http://{engine_ip}:{self.port}",
                 model_names=model_names,
                 added_timestamp=int(time.time()),
+                Id=str(uuid.uuid5(uuid.NAMESPACE_DNS, engine_name)),
                 model_label=model_label,
                 pod_name=engine_name,
                 namespace=self.namespace,

src/vllm_router/services/request_service/request.py

@@ -175,6 +175,11 @@ async def route_general_request(
     request_body = await request.body()
     request_json = await request.json()  # TODO (ApostaC): merge two awaits into one
 
+    if request.query_params:
+        request_endpoint = request.query_params.get("id")
+    else:
+        request_endpoint = None
+
     if hasattr(request.app.state, "callbacks") and (
         response_overwrite := request.app.state.callbacks.pre_request(
             request, request_body, request_json
@@ -206,26 +211,40 @@ async def route_general_request(
     # TODO (ApostaC): merge two awaits into one
     service_discovery = get_service_discovery()
     endpoints = service_discovery.get_endpoint_info()
-    engine_stats = request.app.state.engine_stats_scraper.get_engine_stats()
-    request_stats = request.app.state.request_stats_monitor.get_request_stats(
-        time.time()
-    )
 
     aliases = getattr(service_discovery, "aliases", None)
     if aliases and requested_model in aliases.keys():
         requested_model = aliases[requested_model]
         request_body = replace_model_in_request_body(request_json, requested_model)
         update_content_length(request, request_body)
 
-    # Filter endpoints that have the requested model
-    endpoints = list(filter(lambda x: requested_model in x.model_names, endpoints))
+    if not request_endpoint:
+        endpoints = list(filter(lambda x: requested_model in x.model_names, endpoints))
+        engine_stats = request.app.state.engine_stats_scraper.get_engine_stats()
+        request_stats = request.app.state.request_stats_monitor.get_request_stats(
+            time.time()
+        )
+    else:
+        endpoints = list(
+            filter(
+                lambda x: requested_model in x.model_names and x.Id == request_endpoint,
+                endpoints,
+            )
+        )
+
     if not endpoints:
         return JSONResponse(
             status_code=400, content={"error": f"Model {requested_model} not found."}
         )
 
     logger.debug(f"Routing request {request_id} for model: {requested_model}")
-    if isinstance(request.app.state.router, KvawareRouter) or isinstance(
+    if request_endpoint:
+        server_url = endpoints[0].url
+        logger.debug(
+            f"Routing request {request_id} to engine with Id: {endpoints[0].Id}"
+        )
+
+    elif isinstance(request.app.state.router, KvawareRouter) or isinstance(
         request.app.state.router, PrefixAwareRouter
     ):
         server_url = await request.app.state.router.route_request(
@@ -235,6 +254,7 @@ async def route_general_request(
         server_url = request.app.state.router.route_request(
             endpoints, engine_stats, request_stats, request
         )
+
     curr_time = time.time()
     logger.info(
         f"Routing request {request_id} to {server_url} at {curr_time}, process time = {curr_time - in_router_time:.4f}"