move nixl sidechannel to own entrypoint

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

push missing side channel server

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

Author: wseaton

vllm/entrypoints/nixl_side_channel_server.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import threading
+from typing import Any, Optional
+
+import uvicorn
+from fastapi import FastAPI
+
+from vllm import envs
+from vllm.config import VllmConfig
+from vllm.entrypoints.launcher import serve_http
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class NixlSideChannelServer:
+
+    def __init__(self, vllm_config: VllmConfig, host: str, port: int):
+        self.vllm_config = vllm_config
+        self.host = host
+        self.port = port
+        self.app = FastAPI(title="vLLM NIXL Side Channel Server")
+        self.server = None
+        self.server_thread = None
+        self._setup_routes()
+        
+    def _setup_routes(self):
+        
+        @self.app.get("/get_kv_connector_metadata")
+        @self.app.get("/get_kv_connector_metadata/{dp_rank}")
+        @self.app.get("/get_kv_connector_metadata/{dp_rank}/{tp_rank}")
+        async def get_kv_connector_metadata(dp_rank: Optional[int] = None,
+                                            tp_rank: Optional[int] = None):
+            kv_meta: Optional[dict[str, dict[str, dict[str, Any]]]] = (
+                self.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_meta:
+                    return {}
+                dp_data = kv_meta[dp_rank]
+
+                if tp_rank is not None:
+                    if tp_rank not in dp_data:
+                        return {}
+                    return {dp_rank: {tp_rank: dp_data[tp_rank]}}
+                else:
+                    return {dp_rank: dp_data}
+
+            return kv_meta
+
+    async def start_async(self):
+        if self.server is not None:
+            logger.warning("Side channel server is already running")
+            return
+            
+        logger.info("Starting NIXL side channel server on %s:%s", 
+                   self.host, self.port)
+        
+        # use uvicorn directly to avoid dependency on engine_client
+        config = uvicorn.Config(
+            app=self.app,
+            host=self.host,
+            port=self.port,
+            log_level="info",
+            access_log=True,
+        )
+        self.server = uvicorn.Server(config)
+        
+        # start the server in a background task
+        asyncio.create_task(self.server.serve())
+        logger.info("NIXL side channel server started successfully")
+
+    async def stop_async(self):
+        if self.server is not None:
+            logger.info("Stopping NIXL side channel server")
+            try:
+                self.server.should_exit = True
+                await asyncio.sleep(1)  # give it time to shutdown
+            except Exception as e:
+                logger.warning("Error during side channel server shutdown: %s", e)
+            self.server = None
+            logger.info("NIXL side channel server stopped")
+
+
+def should_start_nixl_side_channel_server(vllm_config: VllmConfig) -> bool:
+    if vllm_config.kv_transfer_config is None:
+        return False
+        
+    return vllm_config.kv_transfer_config.kv_connector == "NixlConnector"
+
+
+async def start_nixl_side_channel_server_if_needed(
+        vllm_config: VllmConfig) -> Optional[NixlSideChannelServer]:
+    if not should_start_nixl_side_channel_server(vllm_config):
+        return None
+        
+    side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
+    side_channel_port = envs.VLLM_NIXL_SIDE_CHANNEL_PORT
+    
+    logger.info("Starting NIXL side channel metadata server on %s:%d",
+               side_channel_host, side_channel_port)
+    
+    server = NixlSideChannelServer(
+        vllm_config, side_channel_host, side_channel_port)
+    await server.start_async()
+    return server

vllm/entrypoints/openai/api_server.py

@@ -93,6 +93,8 @@
 from vllm.entrypoints.openai.serving_transcription import (
     OpenAIServingTranscription, OpenAIServingTranslation)
 from vllm.entrypoints.openai.tool_parsers import ToolParserManager
+from vllm.entrypoints.nixl_side_channel_server import (
+    start_nixl_side_channel_server_if_needed)
 from vllm.entrypoints.utils import (cli_env_setup, load_aware_call,
                                     with_cancellation)
 from vllm.logger import init_logger
@@ -916,32 +918,6 @@ async def show_server_info(raw_request: Request):
         server_info = {"vllm_config": str(raw_request.app.state.vllm_config)}
         return JSONResponse(content=server_info)
 
-    @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: 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_meta:
-                return {}
-            dp_data = kv_meta[dp_rank]
-
-            if tp_rank is not None:
-                if tp_rank not in dp_data:
-                    return {}
-                return {dp_rank: {tp_rank: dp_data[tp_rank]}}
-            else:
-                return {dp_rank: dp_data}
-
-        return kv_meta
 
     @router.post("/reset_prefix_cache")
     async def reset_prefix_cache(raw_request: Request):
@@ -1474,6 +1450,12 @@ async def run_server_worker(listen_address,
         vllm_config = await engine_client.get_vllm_config()
         await init_app_state(engine_client, vllm_config, app.state, args)
 
+        nixl_side_channel_server = None
+        try:
+            nixl_side_channel_server = await start_nixl_side_channel_server_if_needed(vllm_config)
+        except Exception as e:
+            logger.warning("Failed to start NIXL side channel server: %s", e)
+
         logger.info("Starting vLLM API server %d on %s", server_index,
                     listen_address)
         shutdown_task = await serve_http(
@@ -1498,6 +1480,11 @@ async def run_server_worker(listen_address,
     try:
         await shutdown_task
     finally:
+        if nixl_side_channel_server is not None:
+            try:
+                await nixl_side_channel_server.stop_async()
+            except Exception as e:
+                logger.warning("Error stopping NIXL side channel server: %s", e)
         sock.close()