sglang 0.1.15__py3-none-any.whl → 0.1.17__py3-none-any.whl

sglang/srt/server.py

@@ -9,7 +9,8 @@ import os
 import sys
 import threading
 import time
-from typing import List, Optional, Union
+from http import HTTPStatus
+from typing import Optional
 
 # Fix a bug of Python threading
 setattr(threading, "_register_atexit", lambda *args, **kwargs: None)
@@ -20,26 +21,30 @@ import requests
 import uvicorn
 import uvloop
 from fastapi import FastAPI, Request
-from fastapi.responses import Response, StreamingResponse
+from fastapi.responses import JSONResponse, Response, StreamingResponse
 
 from sglang.backend.runtime_endpoint import RuntimeEndpoint
 from sglang.srt.constrained import disable_cache
 from sglang.srt.hf_transformers_utils import get_tokenizer
 from sglang.srt.managers.detokenizer_manager import start_detokenizer_process
 from sglang.srt.managers.io_struct import GenerateReqInput
-from sglang.srt.managers.router.manager import start_router_process
+from sglang.srt.managers.controller.manager_single import start_controller_process as start_controller_process_single
+from sglang.srt.managers.controller.manager_multi import start_controller_process as start_controller_process_multi
 from sglang.srt.managers.tokenizer_manager import TokenizerManager
 from sglang.srt.openai_api_adapter import (
-    v1_completions, v1_chat_completions, load_chat_template_for_openai_api)
-from sglang.srt.server_args import PortArgs, ServerArgs
+    load_chat_template_for_openai_api,
+    v1_chat_completions,
+    v1_completions,
+)
+from sglang.srt.server_args import ModelPortArgs, PortArgs, ServerArgs
 from sglang.srt.utils import (
+    API_KEY_HEADER_NAME,
+    APIKeyValidatorMiddleware,
     allocate_init_ports,
     assert_pkg_version,
     enable_show_time_cost,
-    get_exception_traceback,
-    API_KEY_HEADER_NAME,
-    APIKeyValidatorMiddleware
 )
+from sglang.utils import get_exception_traceback
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
@@ -69,7 +74,7 @@ async def get_server_args():
 
 @app.get("/flush_cache")
 async def flush_cache():
-    await tokenizer_manager.flush_cache()
+    tokenizer_manager.flush_cache()
     return Response(
         content="Cache flushed.\nPlease check backend logs for more details. "
         "(When there are running or waiting requests, the operation will not be performed.)\n",
@@ -77,21 +82,32 @@ async def flush_cache():
     )
 
 
-@app.post("/generate")
-async def generate_request(obj: GenerateReqInput):
-    obj.post_init()
-
+async def generate_request(obj: GenerateReqInput, request: Request):
     if obj.stream:
 
         async def stream_results():
-            async for out in tokenizer_manager.generate_request(obj):
+            try:
+                async for out in tokenizer_manager.generate_request(obj, request):
+                    yield f"data: {json.dumps(out, ensure_ascii=False)}\n\n"
+            except ValueError as e:
+                out = {"error": {"message": str(e)}}
                 yield f"data: {json.dumps(out, ensure_ascii=False)}\n\n"
             yield "data: [DONE]\n\n"
 
-        return StreamingResponse(stream_results(), media_type="text/event-stream")
+        return StreamingResponse(stream_results(), media_type="text/event-stream",
+                                 background=tokenizer_manager.create_abort_task(obj))
+    else:
+        try:
+            ret = await tokenizer_manager.generate_request(obj, request).__anext__()
+            return ret
+        except ValueError as e:
+            return JSONResponse(
+                {"error": {"message": str(e)}}, status_code=HTTPStatus.BAD_REQUEST
+            )
+
 
-    ret = await tokenizer_manager.generate_request(obj).__anext__()
-    return ret
+app.post("/generate")(generate_request)
+app.put("/generate")(generate_request)
 
 
 @app.post("/v1/completions")
@@ -104,7 +120,7 @@ async def openai_v1_chat_completions(raw_request: Request):
     return await v1_chat_completions(tokenizer_manager, raw_request)
 
 
-def launch_server(server_args: ServerArgs, pipe_finish_writer):
+def launch_server(server_args: ServerArgs, pipe_finish_writer, model_overide_args=None):
     global tokenizer_manager
 
     logging.basicConfig(
@@ -126,28 +142,42 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer):
 
     # Allocate ports
     server_args.port, server_args.additional_ports = allocate_init_ports(
-        server_args.port, server_args.additional_ports, server_args.tp_size
+        server_args.port,
+        server_args.additional_ports,
+        server_args.tp_size,
+        server_args.dp_size,
     )
+
+    # Init local models port args
+    ports = server_args.additional_ports
+    tp = server_args.tp_size
+    model_port_args = []
+    for i in range(server_args.dp_size):
+        model_port_args.append(
+            ModelPortArgs(
+                nccl_port=ports[3 + i * (tp + 1)],
+                model_tp_ports=ports[3 + i * (tp + 1) + 1 : 3 + (i + 1) * (tp + 1)],
+            )
+        )
     port_args = PortArgs(
-        tokenizer_port=server_args.additional_ports[0],
-        router_port=server_args.additional_ports[1],
-        detokenizer_port=server_args.additional_ports[2],
-        nccl_port=server_args.additional_ports[3],
-        model_rpc_ports=server_args.additional_ports[4:],
+        tokenizer_port=ports[0],
+        router_port=ports[1],
+        detokenizer_port=ports[2],
+        model_port_args=model_port_args,
     )
 
     # Launch processes
-    tokenizer_manager = TokenizerManager(server_args, port_args)
+    tokenizer_manager = TokenizerManager(server_args, port_args, model_overide_args)
     pipe_router_reader, pipe_router_writer = mp.Pipe(duplex=False)
     pipe_detoken_reader, pipe_detoken_writer = mp.Pipe(duplex=False)
 
+    if server_args.dp_size == 1:
+        start_process = start_controller_process_single
+    else:
+        start_process = start_controller_process_multi
     proc_router = mp.Process(
-        target=start_router_process,
-        args=(
-            server_args,
-            port_args,
-            pipe_router_writer,
-        ),
+        target=start_process,
+        args=(server_args, port_args, pipe_router_writer, model_overide_args),
     )
     proc_router.start()
     proc_detoken = mp.Process(
@@ -167,14 +197,20 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer):
     if router_init_state != "init ok" or detoken_init_state != "init ok":
         proc_router.kill()
         proc_detoken.kill()
-        print(f"Initialization failed. router_init_state: {router_init_state}", flush=True)
-        print(f"Initialization failed. detoken_init_state: {detoken_init_state}", flush=True)
+        print(
+            f"Initialization failed. router_init_state: {router_init_state}", flush=True
+        )
+        print(
+            f"Initialization failed. detoken_init_state: {detoken_init_state}",
+            flush=True,
+        )
         sys.exit(1)
     assert proc_router.is_alive() and proc_detoken.is_alive()
 
     if server_args.api_key and server_args.api_key != "":
         app.add_middleware(APIKeyValidatorMiddleware, api_key=server_args.api_key)
 
+    # Send a warmup request
     def _wait_and_warmup():
         headers = {}
         url = server_args.url()
@@ -192,20 +228,21 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer):
 
         # Send a warmup request
         try:
-            res = requests.post(
-                url + "/generate",
-                json={
-                    "text": "Say this is a warmup request.",
-                    "sampling_params": {
-                        "temperature": 0,
-                        "max_new_tokens": 16,
+            for _ in range(server_args.dp_size):
+                res = requests.post(
+                    url + "/generate",
+                    json={
+                        "text": "The capital city of France is",
+                        "sampling_params": {
+                            "temperature": 0,
+                            "max_new_tokens": 16,
+                        },
                     },
-                },
-                headers=headers,
-                timeout=60,
-            )
-            assert res.status_code == 200
-        except Exception as e:
+                    headers=headers,
+                    timeout=600,
+                )
+                assert res.status_code == 200
+        except Exception:
             if pipe_finish_writer is not None:
                 pipe_finish_writer.send(get_exception_traceback())
             print(f"Initialization failed. warmup error: {e}")
@@ -216,6 +253,8 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer):
 
     t = threading.Thread(target=_wait_and_warmup)
     t.start()
+
+    # Listen for requests
     try:
         uvicorn.run(
             app,
@@ -232,16 +271,21 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer):
 class Runtime:
     def __init__(
         self,
-        log_evel="error",
+        log_level: str = "error",
+        model_overide_args: Optional[dict] = None,
         *args,
         **kwargs,
     ):
         """See the arguments in server_args.py::ServerArgs"""
-        self.server_args = ServerArgs(*args, log_level=log_evel, **kwargs)
+        self.server_args = ServerArgs(*args, log_level=log_level, **kwargs)
 
         # Pre-allocate ports
         self.server_args.port, self.server_args.additional_ports = allocate_init_ports(
-            self.server_args.port, self.server_args.additional_ports, self.server_args.tp_size)
+            self.server_args.port,
+            self.server_args.additional_ports,
+            self.server_args.tp_size,
+            self.server_args.dp_size,
+        )
 
         self.url = self.server_args.url()
         self.generate_url = (
@@ -250,7 +294,10 @@ class Runtime:
 
         self.pid = None
         pipe_reader, pipe_writer = mp.Pipe(duplex=False)
-        proc = mp.Process(target=launch_server, args=(self.server_args, pipe_writer))
+        proc = mp.Process(
+            target=launch_server,
+            args=(self.server_args, pipe_writer, model_overide_args),
+        )
         proc.start()
         pipe_writer.close()
         self.pid = proc.pid
@@ -262,7 +309,9 @@ class Runtime:
 
         if init_state != "init ok":
             self.shutdown()
-            raise RuntimeError("Initialization failed. Please see the error messages above.")
+            raise RuntimeError(
+                "Initialization failed. Please see the error messages above."
+            )
 
         self.endpoint = RuntimeEndpoint(self.url)
 
@@ -314,4 +363,4 @@ class Runtime:
                         pos += len(cur)
 
     def __del__(self):
-        self.shutdown()
+        self.shutdown()

sglang/srt/server_args.py

@@ -2,6 +2,7 @@
 
 import argparse
 import dataclasses
+import random
 from typing import List, Optional, Union
 
 
@@ -15,6 +16,7 @@ class ServerArgs:
     chat_template: Optional[str] = None
     trust_remote_code: bool = True
     context_length: Optional[int] = None
+    quantization: Optional[str] = None
 
     # Port
     host: str = "127.0.0.1"
@@ -23,14 +25,15 @@ class ServerArgs:
 
     # Memory and scheduling
     mem_fraction_static: Optional[float] = None
-    max_prefill_num_token: Optional[int] = None
+    max_prefill_tokens: Optional[int] = None
+    max_running_requests: Optional[int] = None
     schedule_heuristic: str = "lpm"
     schedule_conservativeness: float = 1.0
 
     # Other runtime options
     tp_size: int = 1
     stream_interval: int = 8
-    random_seed: int = 42
+    random_seed: Optional[int] = None
 
     # Logging
     log_level: str = "info"
@@ -42,6 +45,10 @@ class ServerArgs:
     # Other
     api_key: str = ""
 
+    # Data parallelism
+    dp_size: int = 1
+    load_balance_method: str = "round_robin"
+
     # Optimization/debug options
     enable_flashinfer: bool = False
     attention_reduce_in_fp32: bool = False
@@ -66,6 +73,9 @@ class ServerArgs:
         elif self.additional_ports is None:
             self.additional_ports = []
 
+        if self.random_seed is None:
+            self.random_seed = random.randint(0, 1 << 30)
+
     @staticmethod
     def add_cli_args(parser: argparse.ArgumentParser):
         parser.add_argument(
@@ -80,10 +90,12 @@ class ServerArgs:
             default=ServerArgs.tokenizer_path,
             help="The path of the tokenizer.",
         )
-        parser.add_argument("--host", type=str, default=ServerArgs.host,
-                            help="The host of the server.")
-        parser.add_argument("--port", type=int, default=ServerArgs.port,
-                            help="The port of the server.")
+        parser.add_argument(
+            "--host", type=str, default=ServerArgs.host, help="The host of the server."
+        )
+        parser.add_argument(
+            "--port", type=int, default=ServerArgs.port, help="The port of the server."
+        )
         parser.add_argument(
             "--additional-ports",
             type=int,
@@ -133,6 +145,12 @@ class ServerArgs:
             default=ServerArgs.context_length,
             help="The model's maximum context length. Defaults to None (will use the value from the model's config.json instead).",
         )
+        parser.add_argument(
+            "--quantization",
+            type=str,
+            default=ServerArgs.quantization,
+            help="The quantization method.",
+        )
         parser.add_argument(
             "--mem-fraction-static",
             type=float,
@@ -140,16 +158,23 @@ class ServerArgs:
             help="The fraction of the memory used for static allocation (model weights and KV cache memory pool). Use a smaller value if you see out-of-memory errors.",
         )
         parser.add_argument(
-            "--max-prefill-num-token",
+            "--max-prefill-tokens",
             type=int,
-            default=ServerArgs.max_prefill_num_token,
+            default=ServerArgs.max_prefill_tokens,
             help="The maximum number of tokens in a prefill batch. The real bound will be the maximum of this value and the model's maximum context length.",
         )
+        parser.add_argument(
+            "--max-running-requests",
+            type=int,
+            default=ServerArgs.max_running_requests,
+            help="The maximum number of running requests.",
+        )
         parser.add_argument(
             "--schedule-heuristic",
             type=str,
             default=ServerArgs.schedule_heuristic,
-            help="Schudule mode: [lpm, weight, random, fcfs]",
+            choices=["lpm", "random", "fcfs", "dfs-weight"],
+            help="Scheduling Heuristic.",
         )
         parser.add_argument(
             "--schedule-conservativeness",
@@ -209,6 +234,24 @@ class ServerArgs:
             help="Set API key of the server",
         )
 
+        # Data parallelism
+        parser.add_argument(
+            "--dp-size",
+            type=int,
+            default=ServerArgs.dp_size,
+            help="Data parallelism size.",
+        )
+        parser.add_argument(
+            "--load-balance-method",
+            type=str,
+            default=ServerArgs.load_balance_method,
+            help="Load balancing strategy for data parallelism.",
+            choices=[
+                "round_robin",
+                "shortest_queue",
+            ],
+        )
+
         # Optimization/debug options
         parser.add_argument(
             "--enable-flashinfer",
@@ -254,10 +297,15 @@ class ServerArgs:
         )
 
 
+@dataclasses.dataclass
+class ModelPortArgs:
+    nccl_port: int
+    model_tp_ports: List[int]
+
+
 @dataclasses.dataclass
 class PortArgs:
     tokenizer_port: int
     router_port: int
     detokenizer_port: int
-    nccl_port: int
-    model_rpc_ports: List[int]
+    model_port_args: List[ModelPortArgs]