actually use handshake timeout; simplify route

Signed-off-by: Will Eaton <weaton@redhat.com>

Author: wseaton

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

@@ -12,7 +12,8 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional
 from urllib.error import HTTPError, URLError
-from urllib.request import Request, urlopen
+from urllib.request import Request as URLRequest
+from urllib.request import urlopen
 
 import torch
 
@@ -398,7 +399,7 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str):
 
         # Background handshake threads for remote engines
         self._executor = ThreadPoolExecutor(
-            max_workers=4, thread_name_prefix="nixl-handshake")
+            max_workers=1, thread_name_prefix="nixl-handshake")
         # Thread results for handshake completion tracking
         self._handshake_futures: dict[str, Future] = {}
         self._pending_requests: dict[str, list[tuple[str, ReqMeta]]] = {}
@@ -491,13 +492,17 @@ def _nixl_handshake(self, host: str, port: int):
 
         # Use the new endpoint scheme to filter by dp_rank and tp_rank
         # Default to dp_rank 0 and use current tp_rank for optimal filtering
-        url = build_uri("http", host, port, path=f"get_kv_connector_metadata/0/{self.tp_rank}")
+        url = build_uri("http",
+                        host,
+                        port,
+                        path=f"get_kv_connector_metadata/0/{self.tp_rank}")
         logger.debug("Querying metadata on path: %s", url)
 
         try:
-            req = Request(url)
+            req = URLRequest(url)
             logger.debug("About to send HTTP request to %s", url)
-            with urlopen(req, timeout=5.0) as response:
+            with urlopen(req,
+                         timeout=envs.VLLM_NIXL_HANDSHAKE_TIMEOUT) as response:
                 logger.debug("Received HTTP response from %s", url)
                 response_data = response.read().decode('utf-8')
                 res = json.loads(response_data)
@@ -507,27 +512,36 @@ def _nixl_handshake(self, host: str, port: int):
             raise
 
         if res is None:
-            logger.warning("Remote server returned None metadata, skipping handshake")
+            logger.warning(
+                "Remote server returned None metadata, skipping handshake")
             raise RuntimeError("Remote server returned None metadata")
 
-        # With filtered response from new endpoint, we get: {dp_rank: {tp_rank: metadata}}
-        # Since we filtered by dp_rank=0 and tp_rank=self.tp_rank, extract directly
+        # With filtered response from new endpoint, we get:
+        # {dp_rank: {tp_rank: metadata}}
+        # Since we filtered by dp_rank=0 and tp_rank=self.tp_rank,
+        # extract directly.
         if "0" in res and str(self.tp_rank) in res["0"]:
             tp_data = res["0"][str(self.tp_rank)]
             metadata_bytes = tp_data.get("agent_metadata", None)
-            p_remote_rank = self.tp_rank  # Use current tp_rank for filtered response
+            # use current tp_rank for filtered response
+            p_remote_rank = self.tp_rank
         else:
             # Fallback to unfiltered endpoint for heterogeneous TP cases
-            url_fallback = build_uri("http", host, port, path="get_kv_connector_metadata")
-            logger.debug("Using fallback unfiltered endpoint: %s", url_fallback)
-            req = Request(url_fallback)
-            with urlopen(req, timeout=5.0) as response:
+            url_fallback = build_uri("http",
+                                     host,
+                                     port,
+                                     path="get_kv_connector_metadata")
+            logger.debug("Using fallback unfiltered endpoint: %s",
+                         url_fallback)
+            req = URLRequest(url_fallback)
+            with urlopen(req,
+                         timeout=envs.VLLM_NIXL_HANDSHAKE_TIMEOUT) as response:
                 response_data = response.read().decode('utf-8')
                 res = json.loads(response_data)
-            
+
             dp_data = res.get("0", {})
             remote_tp_size = len(dp_data.keys()) if dp_data else 1
-            
+
             # Handle heterogeneous TP mapping
             tp_ratio = self._tp_size[self.engine_id] // remote_tp_size
             p_remote_rank = self.tp_rank // tp_ratio
@@ -952,8 +966,8 @@ def _process_ready_requests(self):
         while True:
             try:
                 req_id, meta = self._ready_requests.get_nowait()
-                logger.debug("Processing ready request %s for engine %s", 
-                            req_id, meta.remote_engine_id)
+                logger.debug("Processing ready request %s for engine %s",
+                             req_id, meta.remote_engine_id)
                 self._read_blocks(
                     request_id=req_id,
                     dst_engine_id=meta.remote_engine_id,
@@ -963,7 +977,7 @@ def _process_ready_requests(self):
                 processed_count += 1
             except queue.Empty:
                 break
-        
+
         if processed_count > 0:
             logger.debug("Processed %d ready requests", processed_count)
 
@@ -972,7 +986,7 @@ def start_load_kv(self, metadata: NixlConnectorMetadata):
         Start loading by triggering non-blocking nixl_xfer.
         We check for these trnxs to complete in each step().
         """
-        
+
         for req_id, meta in metadata.requests.items():
             logger.debug(
                 "start_load_kv for request %s from remote engine %s. "

vllm/entrypoints/openai/api_server.py

@@ -871,28 +871,29 @@ async def show_server_info(raw_request: Request):
     @router.get("/get_kv_connector_metadata")
     @router.get("/get_kv_connector_metadata/{dp_rank}")
     @router.get("/get_kv_connector_metadata/{dp_rank}/{tp_rank}")
-    async def get_kv_connector_metadata(raw_request: Request, dp_rank: int = None, tp_rank: int = None):
-        kv_connector_metadata = raw_request.app.state.vllm_config.cache_config.transfer_handshake_metadata
-        
-        if kv_connector_metadata is None:
-            return JSONResponse(content=None)
-        
-        # Filter by dp_rank if specified
+    async def get_kv_connector_metadata(raw_request: Request,
+                                        dp_rank: Optional[int] = None,
+                                        tp_rank: Optional[int] = None):
+        kv_meta: Optional[dict[str, dict[str, dict[str, Any]]]] = (
+            raw_request.app.state.vllm_config.cache_config.
+            transfer_handshake_metadata)
+
+        if kv_meta is None:
+            return None
+
         if dp_rank is not None:
-            if dp_rank not in kv_connector_metadata:
-                return JSONResponse(content={})
-            dp_data = kv_connector_metadata[dp_rank]
-            
-            # Filter by tp_rank if also specified
+            if dp_rank not in kv_meta:
+                return {}
+            dp_data = kv_meta[dp_rank]
+
             if tp_rank is not None:
                 if tp_rank not in dp_data:
-                    return JSONResponse(content={})
-                return JSONResponse(content={dp_rank: {tp_rank: dp_data[tp_rank]}})
+                    return {}
+                return {dp_rank: {tp_rank: dp_data[tp_rank]}}
             else:
-                return JSONResponse(content={dp_rank: dp_data})
-        
-        # Return all metadata if no filtering
-        return JSONResponse(content=kv_connector_metadata)
+                return {dp_rank: dp_data}
+
+        return kv_meta
 
     @router.post("/reset_prefix_cache")
     async def reset_prefix_cache(raw_request: Request):

vllm/envs.py

@@ -122,6 +122,7 @@
     VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
     VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
     VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5557
+    VLLM_NIXL_HANDSHAKE_TIMEOUT: float = 2.0
     VLLM_ALL2ALL_BACKEND: str = "naive"
     VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840
     VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1
@@ -840,6 +841,11 @@ def get_vllm_port() -> Optional[int]:
     "VLLM_NIXL_SIDE_CHANNEL_PORT":
     lambda: int(os.getenv("VLLM_NIXL_SIDE_CHANNEL_PORT", "5557")),
 
+    # Timeout in seconds for NIXL HTTP handshake requests.
+    # Default is 2 seconds
+    "VLLM_NIXL_HANDSHAKE_TIMEOUT":
+    lambda: float(os.getenv("VLLM_NIXL_HANDSHAKE_TIMEOUT", "2.0")),
+
     # all2all backend for vllm's expert parallel communication
     # Available options:
     # - "naive": naive all2all implementation using all-reduce

vllm/v1/engine/core.py

@@ -453,12 +453,12 @@ def _perform_handshake(
             if hasattr(self, 'transfer_handshake_metadata'
                        ) and self.transfer_handshake_metadata:
                 # self.transfer_handshake_metadata is list of dicts from workers
-                # Each dict already has structure {tp_rank: {dp_rank: metadata}}
+                # Each dict already has structure {dp_rank: {tp_rank: metadata}}
                 # Merge all worker dicts into a single dict
-                content = {}
+                content: dict[str, dict[str, dict[str, Any]]] = {}
                 for worker_dict in self.transfer_handshake_metadata:
                     if worker_dict is not None:
-                        # Deep merge the nested dictionaries instead of overwriting
+                        # Deep merge nested dictionaries instead of overwrite
                         for dp_rank, tp_dict in worker_dict.items():
                             if dp_rank not in content:
                                 content[dp_rank] = {}