From 66fd29740190d639fed1f74623043afb5363c6aa Mon Sep 17 00:00:00 2001 From: dtc Date: Sat, 4 Oct 2025 10:23:20 +0800 Subject: [PATCH] [Misc] improvements for mooncake_connector_v1 (#906) Group contiguous block ids. Add expired time for P node. Copy proxy server from vllm repo. Signed-off-by: Tianchen Ding --- mooncake-integration/CMakeLists.txt | 1 + mooncake-wheel/mooncake/README.md | 2 +- .../mooncake/mooncake_connector_v1.py | 74 ++++- .../mooncake/vllm_v1_proxy_server.py | 272 ++++++++++++++++++ 4 files changed, 335 insertions(+), 14 deletions(-) create mode 100644 mooncake-wheel/mooncake/vllm_v1_proxy_server.py diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index b85ae0ed..5a79ff08 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -82,6 +82,7 @@ install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli_bench.py" "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/transfer_engine_topology_dump.py" "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_connector_v1.py" + "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/vllm_v1_proxy_server.py" DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME} ) diff --git a/mooncake-wheel/mooncake/README.md b/mooncake-wheel/mooncake/README.md index 0caa0627..4e81066e 100644 --- a/mooncake-wheel/mooncake/README.md +++ b/mooncake-wheel/mooncake/README.md @@ -22,7 +22,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct --port 8020 --kv-transfer-config '{"kv_conne #### Proxy ```bash -python {your_vllm_path}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --prefiller-host 192.168.0.2 --prefiller-port 8010 --decoder-host 192.168.0.3 --decoder-port 8020 +python -m mooncake.vllm_v1_proxy_server --prefiller-host 192.168.0.2 --prefiller-port 8010 --decoder-host 192.168.0.3 --decoder-port 8020 ``` Now you can send requests to the proxy server on default port 8000. diff --git a/mooncake-wheel/mooncake/mooncake_connector_v1.py b/mooncake-wheel/mooncake/mooncake_connector_v1.py index d25662a0..fd607460 100644 --- a/mooncake-wheel/mooncake/mooncake_connector_v1.py +++ b/mooncake-wheel/mooncake/mooncake_connector_v1.py @@ -3,11 +3,12 @@ Usage: Adding the following params to vllm command: Prefill: --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer", "kv_connector_module_path":"mooncake.mooncake_connector_v1"}' Decode: --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer", "kv_connector_module_path":"mooncake.mooncake_connector_v1"}' -Proxy: Running tests/v1/kv_connector/nixl_integration/toy_proxy_server.py +Proxy: Running vllm_v1_proxy_server.py """ import contextlib import threading +import time from collections import defaultdict from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor @@ -17,6 +18,7 @@ from os import getenv from typing import TYPE_CHECKING, Any, Optional import msgspec +import numpy as np import torch import zmq @@ -45,6 +47,8 @@ ReqId = str TRANS_DONE = b"trans_done" logger = init_logger(__name__) +VLLM_MOONCAKE_SIDE_CHANNEL_PORT = int(getenv("VLLM_MOONCAKE_SIDE_CHANNEL_PORT", 6557)) +VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT = int(getenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", 120)) class MooncakeAgentMetadata( @@ -70,6 +74,7 @@ class RecvReqMeta: class SendBlockMeta: local_block_ids: list[int] ready: threading.Event + expire_time: float = float("inf") @dataclass @@ -194,7 +199,7 @@ class MooncakeConnectorScheduler: self.engine_id: EngineId = engine_id self.side_channel_host = get_ip() self.side_channel_port = ( - int(getenv("VLLM_MOONCAKE_SIDE_CHANNEL_PORT", 6557)) + + VLLM_MOONCAKE_SIDE_CHANNEL_PORT + vllm_config.parallel_config.data_parallel_rank * vllm_config.parallel_config.tensor_parallel_size) @@ -381,7 +386,7 @@ class MooncakeConnectorWorker: # Mooncake handshake port. self.side_channel_port: int = ( - int(getenv("VLLM_MOONCAKE_SIDE_CHANNEL_PORT", 6557)) + + VLLM_MOONCAKE_SIDE_CHANNEL_PORT + vllm_config.parallel_config.data_parallel_rank * vllm_config.parallel_config.tensor_parallel_size) @@ -486,6 +491,8 @@ class MooncakeConnectorWorker: logger.warning("Request %s not found in reqs_need_send", req_id) return + # Mark it as not expired. We will send it now. + send_meta.expire_time = float("inf") send_reqs.append((req_id, send_meta)) self._send_blocks(send_reqs, meta) @@ -524,15 +531,19 @@ class MooncakeConnectorWorker: if num_local_blocks > num_remote_blocks: local_block_ids = local_block_ids[-num_remote_blocks:] + # Group by indices + group_local_block_ids, group_remote_block_ids = group_concurrent_contiguous( + local_block_ids, remote_block_ids) + for local_layer_addr, remote_layer_addr in zip( local_base_addr, remote_base_addr): - for local_block_id, remote_block_id in zip( - local_block_ids, remote_block_ids): + for group_local_block_id, group_remote_block_id in zip( + group_local_block_ids, group_remote_block_ids): src_ptrs.append(local_layer_addr + - local_block_id * block_len) + group_local_block_id[0] * block_len) dst_ptrs.append(remote_layer_addr + - remote_block_id * block_len) - lengths.append(block_len) + group_remote_block_id[0] * block_len) + lengths.append(block_len * len(group_local_block_id)) logger.debug("Sending kv_caches for request %s (%d blocks) to %s", req_id, num_remote_blocks, remote_session) @@ -629,6 +640,23 @@ class MooncakeConnectorWorker: "and %s requests done recving", self.tp_rank, len(finished_sending_reqs), len(finished_recving_reqs)) + # Handle timeout to avoid stranding blocks on remote. + now = time.perf_counter() + with self.reqs_need_send.lock: + expired_reqs = [ + req_id + for req_id, send_meta in self.reqs_need_send.reqs.items() + if send_meta.expire_time < now + ] + for req_id in expired_reqs: + logger.warning( + "Request %s timed out after %d seconds without " + "being sent. Freeing its blocks on the producer side.", + req_id, VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT) + del self.reqs_need_send.reqs[req_id] + if expired_reqs: + finished_sending_reqs.update(expired_reqs) + return finished_sending_reqs or None, finished_recving_reqs or None def receive_kv(self, path: str, req_blocks: list[tuple[str, list[int]]]): @@ -682,12 +710,14 @@ class MooncakeConnectorWorker: if self.kv_role != "kv_consumer": with self.reqs_need_send.lock: - for req_id, send_meta in metadata.reqs_to_send.items(): - if send_meta: + for req_id, block_ids in metadata.reqs_to_send.items(): + if block_ids: # Already gone through request_finished() - self.reqs_need_send.reqs[ - req_id].local_block_ids = send_meta - self.reqs_need_send.reqs[req_id].ready.set() + send_meta = self.reqs_need_send.reqs[req_id] + send_meta.local_block_ids = block_ids + send_meta.ready.set() + send_meta.expire_time = time.perf_counter( + ) + VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT else: # From update_state_after_alloc(), # but not reach request_finished() yet @@ -712,3 +742,21 @@ def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]: finally: if ctx is not None: ctx.destroy(linger=0) + + +def group_concurrent_contiguous( + src_indices: list[int], + dst_indices: list[int]) -> tuple[list[list[int]], list[list[int]]]: + """Vectorised NumPy implementation.""" + if len(src_indices) == 0: + return [], [] + + brk = np.where((np.diff(src_indices) != 1) + | (np.diff(dst_indices) != 1))[0] + 1 + src_groups = np.split(src_indices, brk) + dst_groups = np.split(dst_indices, brk) + + src_groups = [g.tolist() for g in src_groups] + dst_groups = [g.tolist() for g in dst_groups] + + return src_groups, dst_groups diff --git a/mooncake-wheel/mooncake/vllm_v1_proxy_server.py b/mooncake-wheel/mooncake/vllm_v1_proxy_server.py new file mode 100644 index 00000000..061ae5f7 --- /dev/null +++ b/mooncake-wheel/mooncake/vllm_v1_proxy_server.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# This is a copy from vLLM repo at tests/v1/kv_connector/nixl_integration/toy_proxy_server.py + +import argparse +import itertools +import logging +import os +import uuid +from contextlib import asynccontextmanager + +import httpx +from fastapi import FastAPI, Request +from fastapi.responses import StreamingResponse + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + + +@asynccontextmanager +async def lifespan(app: FastAPI): + """ + Lifespan context manager to handle startup and shutdown events. + """ + # Startup: Initialize client pools for prefiller and decoder services + app.state.prefill_clients = [] + app.state.decode_clients = [] + + # Create prefill clients + for i, (host, port) in enumerate(global_args.prefiller_instances): + prefiller_base_url = f'http://{host}:{port}/v1' + app.state.prefill_clients.append({ + 'client': + httpx.AsyncClient(timeout=None, base_url=prefiller_base_url), + 'host': + host, + 'port': + port, + 'id': + i + }) + + # Create decode clients + for i, (host, port) in enumerate(global_args.decoder_instances): + decoder_base_url = f'http://{host}:{port}/v1' + app.state.decode_clients.append({ + 'client': + httpx.AsyncClient(timeout=None, base_url=decoder_base_url), + 'host': + host, + 'port': + port, + 'id': + i + }) + + # Initialize round-robin iterators + app.state.prefill_iterator = itertools.cycle( + range(len(app.state.prefill_clients))) + app.state.decode_iterator = itertools.cycle( + range(len(app.state.decode_clients))) + + print(f"Initialized {len(app.state.prefill_clients)} prefill clients " + f"and {len(app.state.decode_clients)} decode clients.") + + yield + + # Shutdown: Close all clients + for client_info in app.state.prefill_clients: + await client_info['client'].aclose() + + for client_info in app.state.decode_clients: + await client_info['client'].aclose() + + +# Update FastAPI app initialization to use lifespan +app = FastAPI(lifespan=lifespan) + + +def parse_args(): + parser = argparse.ArgumentParser() + + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--host", type=str, default="localhost") + + # For prefiller instances + parser.add_argument("--prefiller-hosts", + "--prefiller-host", + type=str, + nargs="+", + default=["localhost"]) + parser.add_argument("--prefiller-ports", + "--prefiller-port", + type=int, + nargs="+", + default=[8100]) + + # For decoder instances + parser.add_argument("--decoder-hosts", + "--decoder-host", + type=str, + nargs="+", + default=["localhost"]) + parser.add_argument("--decoder-ports", + "--decoder-port", + type=int, + nargs="+", + default=[8200]) + + args = parser.parse_args() + + # Validate and pair hosts with ports + if len(args.prefiller_hosts) != len(args.prefiller_ports): + raise ValueError( + "Number of prefiller hosts must match number of prefiller ports") + + if len(args.decoder_hosts) != len(args.decoder_ports): + raise ValueError( + "Number of decoder hosts must match number of decoder ports") + + # Create tuples of (host, port) for each service type + args.prefiller_instances = list( + zip(args.prefiller_hosts, args.prefiller_ports)) + args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports)) + + return args + + +def get_next_client(app, service_type: str): + """ + Get the next client in round-robin fashion. + + Args: + app: The FastAPI app instance + service_type: Either 'prefill' or 'decode' + + Returns: + The next client to use + """ + if service_type == 'prefill': + client_idx = next(app.state.prefill_iterator) + return app.state.prefill_clients[client_idx] + elif service_type == 'decode': + client_idx = next(app.state.decode_iterator) + return app.state.decode_clients[client_idx] + else: + raise ValueError(f"Unknown service type: {service_type}") + + +async def send_request_to_service(client_info: dict, endpoint: str, + req_data: dict, request_id: str): + """ + Send a request to a service using a client from the pool. + """ + req_data = req_data.copy() + req_data['kv_transfer_params'] = { + "do_remote_decode": True, + "do_remote_prefill": False, + "remote_engine_id": None, + "remote_block_ids": None, + "remote_host": None, + "remote_port": None + } + req_data["stream"] = False + req_data["max_tokens"] = 1 + if "max_completion_tokens" in req_data: + req_data["max_completion_tokens"] = 1 + if "stream_options" in req_data: + del req_data["stream_options"] + headers = { + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + "X-Request-Id": request_id + } + + response = await client_info['client'].post(endpoint, + json=req_data, + headers=headers) + response.raise_for_status() + + return response + + +async def stream_service_response(client_info: dict, endpoint: str, + req_data: dict, request_id: str): + """ + Asynchronously stream response from a service using a client from the pool. + """ + headers = { + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + "X-Request-Id": request_id + } + + async with client_info['client'].stream("POST", + endpoint, + json=req_data, + headers=headers) as response: + response.raise_for_status() + async for chunk in response.aiter_bytes(): + yield chunk + + +async def _handle_completions(api: str, request: Request): + try: + req_data = await request.json() + request_id = str(uuid.uuid4()) + + # Get the next prefill client in round-robin fashion + prefill_client_info = get_next_client(request.app, 'prefill') + + # Send request to prefill service + response = await send_request_to_service(prefill_client_info, api, + req_data, request_id) + + # Extract the needed fields + response_json = response.json() + kv_transfer_params = response_json.get('kv_transfer_params', {}) + if kv_transfer_params: + req_data["kv_transfer_params"] = kv_transfer_params + + # Get the next decode client in round-robin fashion + decode_client_info = get_next_client(request.app, 'decode') + + logger.debug("Using %s %s", prefill_client_info, decode_client_info) + + # Stream response from decode service + async def generate_stream(): + async for chunk in stream_service_response(decode_client_info, + api, + req_data, + request_id=request_id): + yield chunk + + return StreamingResponse(generate_stream(), + media_type="application/json") + + except Exception as e: + import sys + import traceback + exc_info = sys.exc_info() + print("Error occurred in disagg prefill proxy server" + f" - {api} endpoint") + print(e) + print("".join(traceback.format_exception(*exc_info))) + raise + + +@app.post("/v1/completions") +async def handle_completions(request: Request): + return await _handle_completions("/completions", request) + + +@app.post("/v1/chat/completions") +async def handle_chat_completions(request: Request): + return await _handle_completions("/chat/completions", request) + + +@app.get("/healthcheck") +async def healthcheck(): + """Simple endpoint to check if the server is running.""" + return { + "status": "ok", + "prefill_instances": len(app.state.prefill_clients), + "decode_instances": len(app.state.decode_clients) + } + + +if __name__ == '__main__': + global global_args + global_args = parse_args() + + import uvicorn + uvicorn.run(app, host=global_args.host, port=global_args.port)