sglang 0.5.2rc2__py3-none-any.whl → 0.5.3rc0__py3-none-any.whl

sglang/srt/disaggregation/mini_lb.py

@@ -1,445 +1,6 @@
-"""
-Minimal HTTP load balancer for prefill and decode servers for testing.
-"""
-
-import asyncio
-import dataclasses
-import logging
-import random
-import urllib
-from http import HTTPStatus
-from itertools import chain
-from typing import List, Optional
-
-import aiohttp
-import orjson
-import uvicorn
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import ORJSONResponse, Response, StreamingResponse
-
-from sglang.srt.disaggregation.utils import PDRegistryRequest
-from sglang.srt.utils import maybe_wrap_ipv6_address
-
-AIOHTTP_STREAM_READ_CHUNK_SIZE = (
-    1024 * 64
-)  # 64KB, to prevent aiohttp's "Chunk too big" error
-
-
-def setup_logger():
-    logger = logging.getLogger("pdlb")
-    logger.setLevel(logging.INFO)
-
-    formatter = logging.Formatter(
-        "[PDLB (Python)] %(asctime)s - %(levelname)s - %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-    )
-
-    handler = logging.StreamHandler()
-    handler.setFormatter(formatter)
-    logger.addHandler(handler)
-
-    return logger
-
-
-logger = setup_logger()
-
-
-@dataclasses.dataclass
-class PrefillConfig:
-    url: str
-    bootstrap_port: Optional[int] = None
-
-
-class MiniLoadBalancer:
-    def __init__(
-        self,
-        prefill_configs: List[PrefillConfig],
-        decode_servers: List[str],
-        timeout: int,
-    ):
-        self.prefill_configs = prefill_configs
-        self.prefill_servers = [p.url for p in prefill_configs]
-        self.decode_servers = decode_servers
-        self.timeout = timeout
-
-    def add_prefill_server(self, new_prefill_config: PrefillConfig):
-        self.prefill_configs.append(new_prefill_config)
-        self.prefill_servers.append(new_prefill_config.url)
-
-    def add_decode_server(self, new_decode_server: str):
-        self.decode_servers.append(new_decode_server)
-
-    def select_pair(self):
-        # TODO: return some message instead of panic
-        assert len(self.prefill_configs) > 0, "No prefill servers available"
-        assert len(self.decode_servers) > 0, "No decode servers available"
-
-        prefill_config = random.choice(self.prefill_configs)
-        decode_server = random.choice(self.decode_servers)
-        return prefill_config.url, prefill_config.bootstrap_port, decode_server
-
-    async def generate(
-        self, modified_request, prefill_server, decode_server, endpoint
-    ) -> ORJSONResponse:
-        assert endpoint[0] != "/", f"Endpoint should not start with '/': {endpoint}"
-
-        async with aiohttp.ClientSession(
-            timeout=aiohttp.ClientTimeout(
-                total=self.timeout
-            )  # Add timeout for request reliability
-        ) as session:
-            tasks = [
-                session.post(f"{prefill_server}/{endpoint}", json=modified_request),
-                session.post(f"{decode_server}/{endpoint}", json=modified_request),
-            ]
-
-            # Wait for both responses to complete. Prefill should end first.
-            prefill_response, decode_response = await asyncio.gather(*tasks)
-
-            if "return_logprob" in modified_request:
-
-                prefill_json = await prefill_response.json()
-                ret_json = await decode_response.json()
-
-                # merge `meta_info.input_token_logprobs` from prefill to decode
-                if "meta_info" in ret_json:
-                    if "input_token_logprobs" in ret_json["meta_info"]:
-                        ret_json["meta_info"]["input_token_logprobs"] = (
-                            prefill_json["meta_info"]["input_token_logprobs"]
-                            + ret_json["meta_info"]["input_token_logprobs"]
-                        )
-            else:
-                ret_json = await decode_response.json()
-
-            return ORJSONResponse(
-                content=ret_json,
-                status_code=decode_response.status,
-            )
-
-    async def generate_stream(
-        self, modified_request, prefill_server, decode_server, endpoint="generate"
-    ):
-        assert endpoint[0] != "/", f"Endpoint should not start with '/': {endpoint}"
-
-        async def stream_results():
-            async with aiohttp.ClientSession(
-                timeout=aiohttp.ClientTimeout(
-                    total=self.timeout
-                )  # Add timeout for request reliability
-            ) as session:
-                # Create the tasks for both prefill and decode requests
-                tasks = [
-                    session.post(f"{prefill_server}/{endpoint}", json=modified_request),
-                    session.post(f"{decode_server}/{endpoint}", json=modified_request),
-                ]
-                # Wait for both responses to complete. Since this is streaming, they return immediately.
-                prefill_response, decode_response = await asyncio.gather(*tasks)
-
-                if modified_request.get("return_logprob", False):
-                    prefill_chunks = []
-                    async for chunk in prefill_response.content:
-                        prefill_chunks.append(chunk)
-
-                    first_prefill_chunk = (
-                        prefill_chunks[0].decode("utf-8")[5:].strip("\n")
-                    )
-                    first_prefill_chunk_json = orjson.loads(first_prefill_chunk)
-
-                    async for chunk in decode_response.content:
-                        # Note: This is inefficient
-                        # merge prefill input_token_logprobs, output_token_logprobs to decode
-                        decoded_chunk = chunk.decode("utf-8")
-                        if (
-                            decoded_chunk
-                            and decoded_chunk.startswith("data:")
-                            and "[DONE]" not in decoded_chunk
-                        ):
-                            ret_json = orjson.loads(decoded_chunk[5:].strip("\n"))
-                            ret_json["meta_info"]["input_token_logprobs"] = (
-                                first_prefill_chunk_json["meta_info"][
-                                    "input_token_logprobs"
-                                ]
-                                + ret_json["meta_info"]["input_token_logprobs"]
-                            )
-
-                            yield b"data: " + orjson.dumps(ret_json) + b"\n\n"
-                        else:
-                            yield chunk
-                else:
-                    async for chunk in decode_response.content.iter_chunked(
-                        AIOHTTP_STREAM_READ_CHUNK_SIZE
-                    ):
-                        yield chunk
-
-        return StreamingResponse(
-            stream_results(),
-            media_type="text/event-stream",
-        )
-
-
-app = FastAPI()
-load_balancer: Optional[MiniLoadBalancer] = None
-
-
-@app.get("/health")
-async def health_check():
-    return Response(status_code=200)
-
-
-@app.get("/health_generate")
-async def health_generate():
-    prefill_servers, decode_servers = (
-        load_balancer.prefill_servers,
-        load_balancer.decode_servers,
-    )
-    async with aiohttp.ClientSession() as session:
-        # Create the tasks
-        tasks = []
-        for server in chain(prefill_servers, decode_servers):
-            tasks.append(session.get(f"{server}/health_generate"))
-        for i, response in enumerate(asyncio.as_completed(tasks)):
-            await response
-    return Response(status_code=200)
-
-
-@app.post("/flush_cache")
-async def flush_cache():
-    prefill_servers, decode_servers = (
-        load_balancer.prefill_servers,
-        load_balancer.decode_servers,
-    )
-    async with aiohttp.ClientSession() as session:
-        # Create the tasks
-        tasks = []
-        for server in chain(prefill_servers, decode_servers):
-            tasks.append(session.post(f"{server}/flush_cache"))
-        for i, response in enumerate(asyncio.as_completed(tasks)):
-            await response
-    return Response(status_code=200)
-
-
-@app.get("/get_server_info")
-async def get_server_info():
-    prefill_servers, decode_servers = (
-        load_balancer.prefill_servers,
-        load_balancer.decode_servers,
-    )
-    prefill_infos = []
-    decode_infos = []
-    all_internal_states = []
-
-    async with aiohttp.ClientSession() as session:
-        for server in chain(prefill_servers):
-            server_info = await session.get(f"{server}/get_server_info")
-            prefill_infos.append(await server_info.json())
-        for server in chain(decode_servers):
-            server_info = await session.get(f"{server}/get_server_info")
-            info_json = await server_info.json()
-            decode_infos.append(info_json)
-            # Extract internal_states from decode servers
-            if "internal_states" in info_json:
-                all_internal_states.extend(info_json["internal_states"])
-
-    # Return format expected by bench_one_batch_server.py
-    if all_internal_states:
-        return {
-            "internal_states": all_internal_states,
-            "prefill": prefill_infos,
-            "decode": decode_infos,
-        }
-    else:
-        # Fallback with dummy data if no internal states found
-        return {
-            "internal_states": [
-                {
-                    "last_gen_throughput": 0.0,
-                    "avg_spec_accept_length": None,
-                }
-            ],
-            "prefill": prefill_infos,
-            "decode": decode_infos,
-        }
-
-
-@app.get("/get_model_info")
-async def get_model_info():
-    global load_balancer
-
-    if not load_balancer or not load_balancer.prefill_servers:
-        raise HTTPException(
-            status_code=HTTPStatus.SERVICE_UNAVAILABLE,
-            detail="There is no server registered",
-        )
-
-    target_server_url = load_balancer.prefill_servers[0]
-    endpoint_url = f"{target_server_url}/get_model_info"
-
-    async with aiohttp.ClientSession() as session:
-        try:
-            async with session.get(endpoint_url) as response:
-                if response.status != 200:
-                    error_text = await response.text()
-                    raise HTTPException(
-                        status_code=HTTPStatus.BAD_GATEWAY,
-                        detail=(
-                            f"Failed to get model info from {target_server_url}"
-                            f"Status: {response.status}, Response: {error_text}"
-                        ),
-                    )
-
-                model_info_json = await response.json()
-                return ORJSONResponse(content=model_info_json)
-
-        except aiohttp.ClientError as e:
-            raise HTTPException(
-                status_code=HTTPStatus.SERVICE_UNAVAILABLE,
-                detail=f"Failed to get model info from backend",
-            )
-
-
-@app.post("/generate")
-async def handle_generate_request(request_data: dict):
-    prefill_server, bootstrap_port, decode_server = load_balancer.select_pair()
-
-    # Parse and transform prefill_server for bootstrap data
-    parsed_url = urllib.parse.urlparse(prefill_server)
-    hostname = maybe_wrap_ipv6_address(parsed_url.hostname)
-    modified_request = request_data.copy()
-
-    batch_size = _get_request_batch_size(modified_request)
-    if batch_size is not None:
-        modified_request.update(
-            {
-                "bootstrap_host": [hostname] * batch_size,
-                "bootstrap_port": [bootstrap_port] * batch_size,
-                "bootstrap_room": [
-                    _generate_bootstrap_room() for _ in range(batch_size)
-                ],
-            }
-        )
-    else:
-        modified_request.update(
-            {
-                "bootstrap_host": hostname,
-                "bootstrap_port": bootstrap_port,
-                "bootstrap_room": _generate_bootstrap_room(),
-            }
-        )
-
-    if request_data.get("stream", False):
-        return await load_balancer.generate_stream(
-            modified_request, prefill_server, decode_server, "generate"
-        )
-    else:
-        return await load_balancer.generate(
-            modified_request, prefill_server, decode_server, "generate"
-        )
-
-
-async def _forward_to_backend(request_data: dict, endpoint_name: str):
-    prefill_server, bootstrap_port, decode_server = load_balancer.select_pair()
-
-    # Parse and transform prefill_server for bootstrap data
-    parsed_url = urllib.parse.urlparse(prefill_server)
-    hostname = maybe_wrap_ipv6_address(parsed_url.hostname)
-    modified_request = request_data.copy()
-    modified_request.update(
-        {
-            "bootstrap_host": hostname,
-            "bootstrap_port": bootstrap_port,
-            "bootstrap_room": _generate_bootstrap_room(),
-        }
-    )
-
-    if request_data.get("stream", False):
-        return await load_balancer.generate_stream(
-            modified_request,
-            prefill_server,
-            decode_server,
-            endpoint=endpoint_name,
-        )
-    else:
-        return await load_balancer.generate(
-            modified_request,
-            prefill_server,
-            decode_server,
-            endpoint=endpoint_name,
-        )
-
-
-@app.post("/v1/chat/completions")
-async def handle_chat_completion_request(request_data: dict):
-    return await _forward_to_backend(request_data, "v1/chat/completions")
-
-
-@app.post("/v1/completions")
-async def handle_completion_request(request_data: dict):
-    return await _forward_to_backend(request_data, "v1/completions")
-
-
-def _generate_bootstrap_room():
-    return random.randint(0, 2**63 - 1)
-
-
-# We may utilize `GenerateReqInput`'s logic later
-def _get_request_batch_size(request):
-    if (text := request.get("text")) is not None:
-        return None if isinstance(text, str) else len(text)
-    if (input_ids := request.get("input_ids")) is not None:
-        return None if isinstance(input_ids[0], int) else len(input_ids)
-    return None
-
-
-@app.get("/v1/models")
-async def get_models():
-    prefill_server = load_balancer.prefill_servers[0]  # Get the first prefill server
-    async with aiohttp.ClientSession() as session:
-        try:
-            response = await session.get(f"{prefill_server}/v1/models")
-            if response.status != 200:
-                raise HTTPException(
-                    status_code=response.status,
-                    detail=f"Prefill server error: Status {response.status}",
-                )
-            return ORJSONResponse(content=await response.json())
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-
-@app.post("/register")
-async def register(obj: PDRegistryRequest):
-    if obj.mode == "prefill":
-        load_balancer.add_prefill_server(
-            PrefillConfig(obj.registry_url, obj.bootstrap_port)
-        )
-        logger.info(
-            f"Registered prefill server: {obj.registry_url} with bootstrap port: {obj.bootstrap_port}"
-        )
-    elif obj.mode == "decode":
-        load_balancer.add_decode_server(obj.registry_url)
-        logger.info(f"Registered decode server: {obj.registry_url}")
-    else:
-        raise HTTPException(
-            status_code=400,
-            detail="Invalid mode. Must be either PREFILL or DECODE.",
-        )
-
-    logger.info(
-        f"#Prefill servers: {len(load_balancer.prefill_configs)}, "
-        f"#Decode servers: {len(load_balancer.decode_servers)}"
-    )
-
-    return Response(status_code=200)
-
-
-def run(prefill_configs, decode_addrs, host, port, timeout):
-    global load_balancer
-    load_balancer = MiniLoadBalancer(prefill_configs, decode_addrs, timeout=timeout)
-    uvicorn.run(app, host=host, port=port)
-
-
-if __name__ == "__main__":
-    # FIXME: remove this, use the unified entry point: sglang.srt.disaggregation.launch_lb
-    from sglang.srt.disaggregation.launch_lb import main
-
-    main()
+raise RuntimeError(
+    """The 'mini_lb' module has been relocated to the 'sglang_router' package.
+    We recommend installing 'sglang-router' with Rust support for optimal performance.
+    If you encounter issues building the router with Rust, set the environment variable
+    'SGLANG_ROUTER_BUILD_NO_RUST=1' and add '--mini-lb' to the command line to use the Python version of 'mini_lb'."""
+)

sglang/srt/disaggregation/mooncake/conn.py

@@ -175,6 +175,7 @@ class MooncakeKVManager(BaseKVManager):
         self.disaggregation_mode = disaggregation_mode
         self.init_engine()
         # for p/d multi node infer
+        self.bootstrap_host = server_args.host
         self.bootstrap_port = server_args.disaggregation_bootstrap_port
         self.dist_init_addr = server_args.dist_init_addr
         self.attn_tp_size = get_attention_tp_size()
@@ -458,7 +459,9 @@ class MooncakeKVManager(BaseKVManager):
             dst_head_start_offset = local_tp_rank_in_group * src_heads_per_rank
         else:
             # Send KVCache from 1 prefill instance to multiple decode instances
-            src_head_start_offset = dst_tp_rank_in_group * dst_heads_per_rank
+            src_head_start_offset = (
+                dst_tp_rank_in_group * dst_heads_per_rank
+            ) % src_heads_per_rank
             num_heads_to_send = dst_heads_per_rank
             dst_head_start_offset = 0
 
@@ -1020,6 +1023,7 @@ class MooncakeKVManager(BaseKVManager):
     def _register_to_bootstrap(self):
         """Register KVSender to bootstrap server via HTTP POST."""
         if self.dist_init_addr:
+            # multi node case: bootstrap server's host is dist_init_addr
             if self.dist_init_addr.startswith("["):  # [ipv6]:port or [ipv6]
                 if self.dist_init_addr.endswith("]"):
                     host = self.dist_init_addr
@@ -1028,7 +1032,8 @@ class MooncakeKVManager(BaseKVManager):
             else:
                 host = socket.gethostbyname(self.dist_init_addr.rsplit(":", 1)[0])
         else:
-            host = get_ip()
+            # single node case: bootstrap server's host is same as http server's host
+            host = self.bootstrap_host
             host = maybe_wrap_ipv6_address(host)
 
         bootstrap_server_url = f"{host}:{self.bootstrap_port}"
@@ -1209,7 +1214,7 @@ class MooncakeKVReceiver(BaseKVReceiver):
         mgr: MooncakeKVManager,
         bootstrap_addr: str,
         bootstrap_room: Optional[int] = None,
-        data_parallel_rank: Optional[int] = None,
+        prefill_dp_rank: Optional[int] = None,
     ):
         self.bootstrap_room = bootstrap_room
         self.bootstrap_addr = bootstrap_addr
@@ -1218,7 +1223,6 @@ class MooncakeKVReceiver(BaseKVReceiver):
         self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
         self.conclude_state = None
         self.init_time = None
-        self.data_parallel_rank = data_parallel_rank
 
         if self.bootstrap_addr not in self.kv_mgr.prefill_dp_size_table:
             (
@@ -1317,11 +1321,14 @@ class MooncakeKVReceiver(BaseKVReceiver):
                     self.prefill_attn_tp_size // self.kv_mgr.attn_tp_size
                 ) * (self.prefill_pp_size // self.kv_mgr.pp_size)
 
-        if self.data_parallel_rank is not None:
-            logger.debug(f"Targeting DP rank: {self.data_parallel_rank}")
-            self.target_dp_group = self.data_parallel_rank
+        if prefill_dp_rank is not None:
+            logger.debug(f"Targeting DP rank: {prefill_dp_rank}")
+            self.prefill_dp_rank = prefill_dp_rank
         else:
-            self.target_dp_group = bootstrap_room % self.prefill_dp_size
+            self.prefill_dp_rank = bootstrap_room % self.prefill_dp_size
+
+        # FIXME: alias here: target_dp_group -> prefill_dp_rank
+        self.target_dp_group = self.prefill_dp_rank
 
         self.kv_mgr.required_prefill_response_num_table[self.bootstrap_room] = (
             self.required_prefill_response_num
@@ -1545,7 +1552,8 @@ class MooncakeKVReceiver(BaseKVReceiver):
 
 
 class MooncakeKVBootstrapServer(BaseKVBootstrapServer):
-    def __init__(self, port: int):
+    def __init__(self, host: str, port: int):
+        self.host = host
         self.port = port
         self.app = web.Application()
         self.store = dict()
@@ -1673,7 +1681,7 @@ class MooncakeKVBootstrapServer(BaseKVBootstrapServer):
             self._runner = web.AppRunner(self.app, access_log=access_log)
             self._loop.run_until_complete(self._runner.setup())
 
-            site = web.TCPSite(self._runner, port=self.port)
+            site = web.TCPSite(self._runner, host=self.host, port=self.port)
             self._loop.run_until_complete(site.start())
             self._loop.run_forever()
         except Exception as e: