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

sglang/srt/openai_api_adapter.py

@@ -1,9 +1,12 @@
 """Conversion between OpenAI APIs and native SRT APIs"""
+
+import asyncio
 import json
 import os
+from http import HTTPStatus
 
-from fastapi import HTTPException, Request
-from fastapi.responses import StreamingResponse
+from fastapi import Request
+from fastapi.responses import StreamingResponse, JSONResponse
 
 from sglang.srt.conversation import (
     Conversation,
@@ -26,14 +29,36 @@ from sglang.srt.openai_protocol import (
     CompletionResponseStreamChoice,
     CompletionStreamResponse,
     DeltaMessage,
+    ErrorResponse,
     LogProbs,
     UsageInfo,
 )
-from sglang.srt.utils import jsonify_pydantic_model
-
 
 chat_template_name = None
 
+
+def create_error_response(
+    message: str,
+    err_type: str = "BadRequestError",
+    status_code: HTTPStatus = HTTPStatus.BAD_REQUEST):
+    error = ErrorResponse(message=message,
+                          type=err_type,
+                          code=status_code.value)
+    return JSONResponse(content=error.model_dump(),
+                        status_code=error.code)
+
+
+def create_streaming_error_response(
+    message: str,
+    err_type: str = "BadRequestError",
+    status_code: HTTPStatus = HTTPStatus.BAD_REQUEST) -> str:
+    error = ErrorResponse(message=message,
+                          type=err_type,
+                          code=status_code.value)
+    json_str = json.dumps({"error": error.model_dump()})
+    return json_str
+
+
 def load_chat_template_for_openai_api(chat_template_arg):
     global chat_template_name
 
@@ -73,8 +98,8 @@ async def v1_completions(tokenizer_manager, raw_request: Request):
     request_json = await raw_request.json()
     request = CompletionRequest(**request_json)
 
-    # TODO: Validate the request and return HTTPStatus.BAD_REQUEST if invalid.
-    assert request.n == 1
+    if request.n != 1:
+        return create_error_response("n != 1 is not supported")
 
     adapted_request = GenerateReqInput(
         text=request.prompt,
@@ -92,79 +117,88 @@ async def v1_completions(tokenizer_manager, raw_request: Request):
         return_text_in_logprobs=True,
         stream=request.stream,
     )
-    adapted_request.post_init()
 
     if adapted_request.stream:
 
         async def generate_stream_resp():
             stream_buffer = ""
             n_prev_token = 0
-            async for content in tokenizer_manager.generate_request(adapted_request):
-                text = content["text"]
-                prompt_tokens = content["meta_info"]["prompt_tokens"]
-                completion_tokens = content["meta_info"]["completion_tokens"]
-
-                if not stream_buffer:  # The first chunk
-                    if request.echo:
-                        # Prepend prompt in response text.
-                        text = request.prompt + text
-
-                if request.logprobs:
-                    # The first chunk and echo is enabled.
-                    if not stream_buffer and request.echo:
-                        prefill_token_logprobs = content["meta_info"][
-                            "prefill_token_logprobs"
-                        ]
-                        prefill_top_logprobs = content["meta_info"][
-                            "prefill_top_logprobs"
-                        ]
+            try:
+                async for content in tokenizer_manager.generate_request(
+                    adapted_request, raw_request):
+                    text = content["text"]
+                    prompt_tokens = content["meta_info"]["prompt_tokens"]
+                    completion_tokens = content["meta_info"]["completion_tokens"]
+
+                    if not stream_buffer:  # The first chunk
+                        if request.echo:
+                            # Prepend prompt in response text.
+                            text = request.prompt + text
+
+                    if request.logprobs:
+                        # The first chunk and echo is enabled.
+                        if not stream_buffer and request.echo:
+                            prefill_token_logprobs = content["meta_info"][
+                                "prefill_token_logprobs"
+                            ]
+                            prefill_top_logprobs = content["meta_info"][
+                                "prefill_top_logprobs"
+                            ]
+                        else:
+                            prefill_token_logprobs = None
+                            prefill_top_logprobs = None
+
+                        logprobs = to_openai_style_logprobs(
+                            prefill_token_logprobs=prefill_token_logprobs,
+                            prefill_top_logprobs=prefill_top_logprobs,
+                            decode_token_logprobs=content["meta_info"][
+                                "decode_token_logprobs"
+                            ][n_prev_token:],
+                            decode_top_logprobs=content["meta_info"]["decode_top_logprobs"][
+                                n_prev_token:
+                            ],
+                        )
+
+                        n_prev_token = len(content["meta_info"]["decode_token_logprobs"])
                     else:
-                        prefill_token_logprobs = None
-                        prefill_top_logprobs = None
-
-                    logprobs = to_openai_style_logprobs(
-                        prefill_token_logprobs=prefill_token_logprobs,
-                        prefill_top_logprobs=prefill_top_logprobs,
-                        decode_token_logprobs=content["meta_info"][
-                            "decode_token_logprobs"
-                        ][n_prev_token:],
-                        decode_top_logprobs=content["meta_info"]["decode_top_logprobs"][
-                            n_prev_token:
-                        ],
-                    )
+                        logprobs = None
 
-                    n_prev_token = len(content["meta_info"]["decode_token_logprobs"])
-                else:
-                    logprobs = None
-
-                delta = text[len(stream_buffer) :]
-                stream_buffer = content["text"]
-                choice_data = CompletionResponseStreamChoice(
-                    index=0,
-                    text=delta,
-                    logprobs=logprobs,
-                    finish_reason=None,
-                )
-                chunk = CompletionStreamResponse(
-                    id=content["meta_info"]["id"],
-                    object="text_completion",
-                    choices=[choice_data],
-                    model=request.model,
-                    usage=UsageInfo(
-                        prompt_tokens=prompt_tokens,
-                        completion_tokens=completion_tokens,
-                        total_tokens=prompt_tokens + completion_tokens,
-                    ),
-                )
-                yield f"data: {jsonify_pydantic_model(chunk)}\n\n"
+                    delta = text[len(stream_buffer) :]
+                    stream_buffer = content["text"]
+                    choice_data = CompletionResponseStreamChoice(
+                        index=0,
+                        text=delta,
+                        logprobs=logprobs,
+                        finish_reason=content["meta_info"]["finish_reason"],
+                    )
+                    chunk = CompletionStreamResponse(
+                        id=content["meta_info"]["id"],
+                        object="text_completion",
+                        choices=[choice_data],
+                        model=request.model,
+                        usage=UsageInfo(
+                            prompt_tokens=prompt_tokens,
+                            completion_tokens=completion_tokens,
+                            total_tokens=prompt_tokens + completion_tokens,
+                        ),
+                    )
+                    yield f"data: {chunk.model_dump_json()}\n\n"
+            except ValueError as e:
+                error = create_streaming_error_response(str(e))
+                yield f"data: {error}\n\n"
             yield "data: [DONE]\n\n"
 
-        return StreamingResponse(generate_stream_resp(), media_type="text/event-stream")
+        return StreamingResponse(generate_stream_resp(), media_type="text/event-stream",
+                                 background=tokenizer_manager.create_abort_task(adapted_request))
 
     # Non-streaming response.
-    ret = await tokenizer_manager.generate_request(adapted_request).__anext__()
-    ret = ret[0] if isinstance(ret, list) else ret
+    try:
+        ret = await tokenizer_manager.generate_request(
+            adapted_request, raw_request).__anext__()
+    except ValueError as e:
+        return create_error_response(str(e))
 
+    ret = ret[0] if isinstance(ret, list) else ret
     prompt_tokens = ret["meta_info"]["prompt_tokens"]
     completion_tokens = ret["meta_info"]["completion_tokens"]
     text = ret["text"]
@@ -192,7 +226,7 @@ async def v1_completions(tokenizer_manager, raw_request: Request):
         index=0,
         text=text,
         logprobs=logprobs,
-        finish_reason=None,  # TODO(comaniac): Add finish reason.
+        finish_reason=ret["meta_info"]["finish_reason"],
     )
     response = CompletionResponse(
         id=ret["meta_info"]["id"],
@@ -211,8 +245,8 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
     request_json = await raw_request.json()
     request = ChatCompletionRequest(**request_json)
 
-    # TODO: Validate the request and return HTTPStatus.BAD_REQUEST if invalid.
-    assert request.n == 1
+    if request.n != 1:
+        return create_error_response("n != 1 is not supported")
 
     # Prep the data needed for the underlying GenerateReqInput:
     #  - prompt: The full prompt string.
@@ -257,7 +291,6 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
         },
         stream=request.stream,
     )
-    adapted_request.post_init()
 
     if adapted_request.stream:
 
@@ -265,46 +298,58 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
             is_first = True
 
             stream_buffer = ""
-            async for content in tokenizer_manager.generate_request(adapted_request):
-                if is_first:
-                    # First chunk with role
-                    is_first = False
+            try:
+                async for content in tokenizer_manager.generate_request(adapted_request, raw_request):
+                    if is_first:
+                        # First chunk with role
+                        is_first = False
+                        choice_data = ChatCompletionResponseStreamChoice(
+                            index=0,
+                            delta=DeltaMessage(role="assistant"),
+                            finish_reason=content["meta_info"]["finish_reason"],
+                        )
+                        chunk = ChatCompletionStreamResponse(
+                            id=content["meta_info"]["id"],
+                            choices=[choice_data],
+                            model=request.model,
+                        )
+                        yield f"data: {chunk.model_dump_json()}\n\n"
+
+                    text = content["text"]
+                    delta = text[len(stream_buffer) :]
+                    stream_buffer = text
                     choice_data = ChatCompletionResponseStreamChoice(
                         index=0,
-                        delta=DeltaMessage(role="assistant"),
-                        finish_reason=None,
+                        delta=DeltaMessage(content=delta),
+                        finish_reason=content["meta_info"]["finish_reason"],
                     )
                     chunk = ChatCompletionStreamResponse(
                         id=content["meta_info"]["id"],
                         choices=[choice_data],
                         model=request.model,
                     )
-                    yield f"data: {jsonify_pydantic_model(chunk)}\n\n"
-
-                text = content["text"]
-                delta = text[len(stream_buffer) :]
-                stream_buffer = text
-                choice_data = ChatCompletionResponseStreamChoice(
-                    index=0, delta=DeltaMessage(content=delta), finish_reason=None
-                )
-                chunk = ChatCompletionStreamResponse(
-                    id=content["meta_info"]["id"],
-                    choices=[choice_data],
-                    model=request.model,
-                )
-                yield f"data: {jsonify_pydantic_model(chunk)}\n\n"
+                    yield f"data: {chunk.model_dump_json()}\n\n"
+            except ValueError as e:
+                error = create_streaming_error_response(str(e))
+                yield f"data: {error}\n\n"
             yield "data: [DONE]\n\n"
 
-        return StreamingResponse(generate_stream_resp(), media_type="text/event-stream")
+        return StreamingResponse(generate_stream_resp(), media_type="text/event-stream",
+                                 background=tokenizer_manager.create_abort_task(adapted_request))
 
     # Non-streaming response.
-    ret = await tokenizer_manager.generate_request(adapted_request).__anext__()
+    try:
+        ret = await tokenizer_manager.generate_request(
+            adapted_request, raw_request).__anext__()
+    except ValueError as e:
+        return create_error_response(str(e))
+
     prompt_tokens = ret["meta_info"]["prompt_tokens"]
     completion_tokens = ret["meta_info"]["completion_tokens"]
     choice_data = ChatCompletionResponseChoice(
         index=0,
         message=ChatMessage(role="assistant", content=ret["text"]),
-        finish_reason=None,  # TODO(comaniac): Add finish reason.
+        finish_reason=ret["meta_info"]["finish_reason"],
     )
     response = ChatCompletionResponse(
         id=ret["meta_info"]["id"],
@@ -332,7 +377,7 @@ def to_openai_style_logprobs(
             ret_logprobs.tokens.append(token_text)
             ret_logprobs.token_logprobs.append(logprob)
 
-            # Not Supported yet
+            # Not supported yet
             ret_logprobs.text_offset.append(-1)
 
     def append_top_logprobs(top_logprobs):
@@ -353,4 +398,4 @@ def to_openai_style_logprobs(
     if decode_top_logprobs is not None:
         append_top_logprobs(decode_top_logprobs)
 
-    return ret_logprobs
+    return ret_logprobs

sglang/srt/openai_protocol.py

@@ -1,4 +1,5 @@
 """pydantic models for OpenAI API protocol"""
+
 import time
 from typing import Dict, List, Optional, Union
 
@@ -6,6 +7,14 @@ from pydantic import BaseModel, Field
 from typing_extensions import Literal
 
 
+class ErrorResponse(BaseModel):
+    object: str = "error"
+    message: str
+    type: str
+    param: Optional[str] = None
+    code: int
+
+
 class LogProbs(BaseModel):
     text_offset: List[int] = Field(default_factory=list)
     token_logprobs: List[Optional[float]] = Field(default_factory=list)
@@ -178,4 +187,4 @@ class ChatCompletionStreamResponse(BaseModel):
     object: str = "chat.completion.chunk"
     created: int = Field(default_factory=lambda: int(time.time()))
     model: str
-    choices: List[ChatCompletionResponseStreamChoice]
+    choices: List[ChatCompletionResponseStreamChoice]

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)
@@ -27,19 +28,23 @@ 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,24 +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
+            )
+
 
-    try:
-        ret = await tokenizer_manager.generate_request(obj).__anext__()
-        return ret
-    except ValueError as e:
-        return JSONResponse({"error": str(e)}, status_code=400)
+app.post("/generate")(generate_request)
+app.put("/generate")(generate_request)
 
 
 @app.post("/v1/completions")
@@ -129,14 +142,28 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer, model_overide_arg
 
     # 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
@@ -144,8 +171,12 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer, model_overide_arg
     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,
+        target=start_process,
         args=(server_args, port_args, pipe_router_writer, model_overide_args),
     )
     proc_router.start()
@@ -179,6 +210,7 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer, model_overide_arg
     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()
@@ -190,27 +222,27 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer, model_overide_arg
             time.sleep(0.5)
             try:
                 requests.get(url + "/get_model_info", timeout=5, headers=headers)
-                success = True  # Set flag to True if request succeeds
                 break
             except requests.exceptions.RequestException as e:
                 pass
 
         # 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=600,
-            )
-            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}")
@@ -221,6 +253,8 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer, model_overide_arg
 
     t = threading.Thread(target=_wait_and_warmup)
     t.start()
+
+    # Listen for requests
     try:
         uvicorn.run(
             app,
@@ -237,19 +271,20 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer, model_overide_arg
 class Runtime:
     def __init__(
         self,
-        log_evel: str = "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.dp_size,
         )
 
         self.url = self.server_args.url()