Adding check for engine sleeping status at service_discovery

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

Author: dumb0002

src/vllm_router/service_discovery.py

@@ -49,6 +49,9 @@ class EndpointInfo:
     # Model label
     model_label: str
 
+    # Endpoint's sleep status
+    sleep: bool
+
 
 class ServiceDiscovery(metaclass=abc.ABCMeta):
     @abc.abstractmethod
@@ -103,7 +106,7 @@ def get_endpoint_info(self) -> List[EndpointInfo]:
         """
 
         endpoint_infos = [
-            EndpointInfo(url, model, self.added_timestamp, model_label)
+            EndpointInfo(url, model, self.added_timestamp, model_label, False)
             for url, model, model_label in zip(
                 self.urls, self.models, self.model_labels
             )
@@ -157,6 +160,33 @@ def _check_pod_ready(container_statuses):
         ready_count = sum(1 for status in container_statuses if status.ready)
         return ready_count == len(container_statuses)
 
+    def _get_engine_sleep_status(self, pod_ip) -> Optional[bool]:
+        """
+        Get the engine sleeping status by querying the engine's
+        '/is_sleeping' endpoint.
+
+        Args:
+            pod_ip: the IP address of the pod running the engine
+
+        Returns:
+            the sleep status of the engine
+        """
+        url = f"http://{pod_ip}:{self.port}/is_sleeping"
+        sleep = False
+        try:
+            headers = None
+            if VLLM_API_KEY := os.getenv("VLLM_API_KEY"):
+                logger.info(f"Using vllm server authentication")
+                headers = {"Authorization": f"Bearer {VLLM_API_KEY}"}
+            response = requests.get(url, headers=headers)
+            response.raise_for_status()
+            sleep = response.json()["is_sleeping"]
+        except Exception as e:
+            logger.error(f"Failed to get the sleep status for engine at {url}: {e}")
+            return None
+
+        return sleep
+
     def _get_model_name(self, pod_ip) -> Optional[str]:
         """
         Get the model name of the serving engine pod by querying the pod's
@@ -238,12 +268,14 @@ def _add_engine(
             f"Discovered new serving engine {engine_name} at "
             f"{engine_ip}, running model: {model_name}"
         )
+
         with self.available_engines_lock:
             self.available_engines[engine_name] = EndpointInfo(
                 url=f"http://{engine_ip}:{self.port}",
                 model_name=model_name,
                 added_timestamp=int(time.time()),
                 model_label=model_label,
+                sleep=self._get_engine_sleep_status(engine_ip),
             )
 
     def _delete_engine(self, engine_name: str):

src/vllm_router/services/request_service/request.py

@@ -217,7 +217,11 @@ async def route_general_request(
         request_body = replace_model_in_request_body(request_json, requested_model)
         update_content_length(request, request_body)
 
-    endpoints = list(filter(lambda x: x.model_name == requested_model, endpoints))
+    endpoints = list(
+        filter(
+            lambda x: x.model_name == requested_model and x.sleep == False, endpoints
+        )
+    )
     if not endpoints:
         return JSONResponse(
             status_code=400, content={"error": f"Model {requested_model} not found."}
@@ -286,7 +290,11 @@ async def route_disaggregated_prefill_request(
         time.time()
     )
 
-    endpoints = list(filter(lambda x: x.model_name == requested_model, endpoints))
+    endpoints = list(
+        filter(
+            lambda x: x.model_name == requested_model and x.sleep == False, endpoints
+        )
+    )
     if not endpoints:
         return JSONResponse(
             status_code=400, content={"error": f"Model {requested_model} not found."}