sglang 0.3.6.post2__py3-none-any.whl → 0.4.0__py3-none-any.whl

sglang/srt/server.py

@@ -23,6 +23,7 @@ import json
 import logging
 import multiprocessing as mp
 import os
+import signal
 import threading
 import time
 from http import HTTPStatus
@@ -51,8 +52,11 @@ from sglang.srt.managers.io_struct import (
     CloseSessionReqInput,
     EmbeddingReqInput,
     GenerateReqInput,
+    GetWeightsByNameReqInput,
+    InitWeightsUpdateGroupReqInput,
     OpenSessionReqInput,
-    UpdateWeightReqInput,
+    UpdateWeightFromDiskReqInput,
+    UpdateWeightsFromDistributedReqInput,
 )
 from sglang.srt.managers.scheduler import run_scheduler_process
 from sglang.srt.managers.tokenizer_manager import TokenizerManager
@@ -79,7 +83,7 @@ from sglang.srt.utils import (
     configure_logger,
     delete_directory,
     is_port_available,
-    kill_child_process,
+    kill_process_tree,
     maybe_set_triton_cache_manager,
     prepare_model_and_tokenizer,
     set_prometheus_multiproc_dir,
@@ -92,7 +96,7 @@ logger = logging.getLogger(__name__)
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
-
+# Fast API
 app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
@@ -103,7 +107,7 @@ app.add_middleware(
 )
 
 tokenizer_manager: TokenizerManager = None
-_max_total_num_tokens = None
+scheduler_info: Dict = None
 
 ##### Native API endpoints #####
 
@@ -149,13 +153,11 @@ async def get_model_info():
 
 @app.get("/get_server_info")
 async def get_server_info():
-    try:
-        return await _get_server_info()
-
-    except Exception as e:
-        return ORJSONResponse(
-            {"error": {"message": str(e)}}, status_code=HTTPStatus.BAD_REQUEST
-        )
+    return {
+        **dataclasses.asdict(tokenizer_manager.server_args),  # server args
+        **scheduler_info,
+        "version": __version__,
+    }
 
 
 @app.post("/flush_cache")
@@ -171,7 +173,7 @@ async def flush_cache():
 
 @app.get("/start_profile")
 @app.post("/start_profile")
-async def start_profile():
+async def start_profile_async():
     """Start profiling."""
     tokenizer_manager.start_profile()
     return Response(
@@ -182,7 +184,7 @@ async def start_profile():
 
 @app.get("/stop_profile")
 @app.post("/stop_profile")
-async def stop_profile():
+async def stop_profile_async():
     """Stop profiling."""
     tokenizer_manager.stop_profile()
     return Response(
@@ -191,11 +193,11 @@ async def stop_profile():
     )
 
 
-@app.post("/update_weights")
+@app.post("/update_weights_from_disk")
 @time_func_latency
-async def update_weights(obj: UpdateWeightReqInput, request: Request):
-    """Update the weights inplace without re-launching the server."""
-    success, message = await tokenizer_manager.update_weights(obj, request)
+async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: Request):
+    """Update the weights from disk inplace without re-launching the server."""
+    success, message = await tokenizer_manager.update_weights_from_disk(obj, request)
     content = {"success": success, "message": message}
     if success:
         return ORJSONResponse(
@@ -209,6 +211,52 @@ async def update_weights(obj: UpdateWeightReqInput, request: Request):
         )
 
 
+@app.post("/init_weights_update_group")
+async def init_weights_update_group(
+    obj: InitWeightsUpdateGroupReqInput, request: Request
+):
+    """Initialize the parameter update group."""
+    success, message = await tokenizer_manager.init_weights_update_group(obj, request)
+    content = {"success": success, "message": message}
+    if success:
+        return ORJSONResponse(content, status_code=200)
+    else:
+        return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)
+
+
+@app.post("/update_weights_from_distributed")
+async def update_weights_from_distributed(
+    obj: UpdateWeightsFromDistributedReqInput, request: Request
+):
+    """Update model parameter from distributed online."""
+    success, message = await tokenizer_manager.update_weights_from_distributed(
+        obj, request
+    )
+    content = {"success": success, "message": message}
+    if success:
+        return ORJSONResponse(content, status_code=200)
+    else:
+        return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)
+
+
+@app.api_route("/get_weights_by_name", methods=["GET", "POST"])
+async def get_weights_by_name(obj: GetWeightsByNameReqInput, request: Request):
+    """Get model parameter by name."""
+    try:
+        ret = await tokenizer_manager.get_weights_by_name(obj, request)
+        if ret is None:
+            return ORJSONResponse(
+                {"error": {"message": "Get parameter by name failed"}},
+                status_code=HTTPStatus.BAD_REQUEST,
+            )
+        else:
+            return ORJSONResponse(ret, status_code=200)
+    except Exception as e:
+        return ORJSONResponse(
+            {"error": {"message": str(e)}}, status_code=HTTPStatus.BAD_REQUEST
+        )
+
+
 @app.api_route("/open_session", methods=["GET", "POST"])
 async def open_session(obj: OpenSessionReqInput, request: Request):
     """Open a session, and return its unique session id."""
@@ -233,6 +281,8 @@ async def close_session(obj: CloseSessionReqInput, request: Request):
         )
 
 
+# fastapi implicitly converts json in the request to obj (dataclass)
+@app.api_route("/generate", methods=["POST", "PUT"])
 @time_func_latency
 async def generate_request(obj: GenerateReqInput, request: Request):
     """Handle a generate request."""
@@ -266,11 +316,7 @@ async def generate_request(obj: GenerateReqInput, request: Request):
             )
 
 
-# fastapi implicitly converts json in the request to obj (dataclass)
-app.post("/generate")(generate_request)
-app.put("/generate")(generate_request)
-
-
+@app.api_route("/encode", methods=["POST", "PUT"])
 @time_func_latency
 async def encode_request(obj: EmbeddingReqInput, request: Request):
     """Handle an embedding request."""
@@ -283,10 +329,7 @@ async def encode_request(obj: EmbeddingReqInput, request: Request):
         )
 
 
-app.post("/encode")(encode_request)
-app.put("/encode")(encode_request)
-
-
+@app.api_route("/encode", methods=["POST", "PUT"])
 @time_func_latency
 async def classify_request(obj: EmbeddingReqInput, request: Request):
     """Handle a reward model request. Now the arguments and return values are the same as embedding models."""
@@ -299,10 +342,6 @@ async def classify_request(obj: EmbeddingReqInput, request: Request):
         )
 
 
-app.post("/classify")(classify_request)
-app.put("/classify")(classify_request)
-
-
 ##### OpenAI-compatible API endpoints #####
 
 
@@ -380,11 +419,11 @@ def launch_engine(
     server_args: ServerArgs,
 ):
     """
-    Launch the Tokenizer Manager in the main process, the Scheduler in a subprocess, and the Detokenizer Manager in another subprocess.
+    Launch the TokenizerManager in the main process, the Scheduler in a subprocess, and the DetokenizerManager in another subprocess.
     """
 
     global tokenizer_manager
-    global _max_total_num_tokens
+    global scheduler_info
 
     # Configure global environment
     configure_logger(server_args)
@@ -450,8 +489,8 @@ def launch_engine(
     if server_args.chat_template:
         load_chat_template_for_openai_api(tokenizer_manager, server_args.chat_template)
 
-    # Wait for model to finish loading & get max token nums
-    scheduler_info = []
+    # Wait for model to finish loading
+    scheduler_infos = []
     for i in range(len(scheduler_pipe_readers)):
         data = scheduler_pipe_readers[i].recv()
 
@@ -459,10 +498,10 @@ def launch_engine(
             raise RuntimeError(
                 "Initialization failed. Please see the error messages above."
             )
-        scheduler_info.append(data)
+        scheduler_infos.append(data)
 
     # Assume all schedulers have same max_total_num_tokens
-    _max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"]
+    scheduler_info = scheduler_infos[0]
 
 
 def launch_server(
@@ -476,12 +515,12 @@ def launch_server(
 
     1. HTTP server: A FastAPI server that routes requests to the engine.
     2. SRT engine:
-        1. Tokenizer Manager: Tokenizes the requests and sends them to the scheduler.
+        1. TokenizerManager: Tokenizes the requests and sends them to the scheduler.
         2. Scheduler (subprocess): Receives requests from the Tokenizer Manager, schedules batches, forwards them, and sends the output tokens to the Detokenizer Manager.
-        3. Detokenizer Manager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager.
+        3. DetokenizerManager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager.
 
     Note:
-    1. The HTTP server and Tokenizer Manager both run in the main process.
+    1. The HTTP server and TokenizerManager both run in the main process.
     2. Inter-process communication is done through ICP (each process uses a different port) via the ZMQ library.
     """
     launch_engine(server_args=server_args)
@@ -490,7 +529,7 @@ def launch_server(
     if server_args.api_key:
         add_api_key_middleware(app, server_args.api_key)
 
-    # add prometheus middleware
+    # Add prometheus middleware
     if server_args.enable_metrics:
         add_prometheus_middleware(app)
         enable_func_timer()
@@ -502,7 +541,7 @@ def launch_server(
     t.start()
 
     try:
-        # Listen for HTTP requests
+        # Update logging configs
         LOGGING_CONFIG["formatters"]["default"][
             "fmt"
         ] = "[%(asctime)s] %(levelprefix)s %(message)s"
@@ -511,6 +550,8 @@ def launch_server(
             "fmt"
         ] = '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s'
         LOGGING_CONFIG["formatters"]["access"]["datefmt"] = "%Y-%m-%d %H:%M:%S"
+
+        # Listen for HTTP requests
         uvicorn.run(
             app,
             host=server_args.host,
@@ -523,15 +564,6 @@ def launch_server(
         t.join()
 
 
-async def _get_server_info():
-    return {
-        **dataclasses.asdict(tokenizer_manager.server_args),  # server args
-        "memory_pool_size": await tokenizer_manager.get_memory_pool_size(),  # memory pool size
-        "max_total_num_tokens": _max_total_num_tokens,  # max total num tokens
-        "version": __version__,
-    }
-
-
 def _set_envs_and_config(server_args: ServerArgs):
     # Set global environments
     os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -562,6 +594,15 @@ def _set_envs_and_config(server_args: ServerArgs):
             "at https://docs.flashinfer.ai/installation.html.",
         )
 
+    # Register the signal handler.
+    # The child processes will send SIGQUIT to this process when any error happens
+    # This process then clean up the whole process tree
+    def sigquit_handler(signum, frame):
+        kill_process_tree(os.getpid())
+
+    signal.signal(signal.SIGQUIT, sigquit_handler)
+
+    # Set mp start method
     mp.set_start_method("spawn", force=True)
 
 
@@ -588,7 +629,7 @@ def _wait_and_warmup(server_args, pipe_finish_writer):
         if pipe_finish_writer is not None:
             pipe_finish_writer.send(last_traceback)
         logger.error(f"Initialization failed. warmup error: {last_traceback}")
-        kill_child_process(include_self=True)
+        kill_process_tree(os.getpid())
         return
 
     model_info = res.json()
@@ -621,9 +662,10 @@ def _wait_and_warmup(server_args, pipe_finish_writer):
         if pipe_finish_writer is not None:
             pipe_finish_writer.send(last_traceback)
         logger.error(f"Initialization failed. warmup error: {last_traceback}")
-        kill_child_process(include_self=True)
+        kill_process_tree(os.getpid())
         return
 
+    # Debug print
     # logger.info(f"{res.json()=}")
 
     logger.info("The server is fired up and ready to roll!")
@@ -634,11 +676,218 @@ def _wait_and_warmup(server_args, pipe_finish_writer):
         delete_directory(server_args.model_path)
 
 
+STREAM_END_SYMBOL = b"data: [DONE]"
+STREAM_CHUNK_START_SYMBOL = b"data:"
+
+
+class Engine:
+    """
+    SRT Engine without an HTTP server layer.
+
+    This class provides a direct inference engine without the need for an HTTP server. It is designed for use cases where
+    launching the HTTP server adds unnecessary complexity or overhead,
+    """
+
+    def __init__(self, log_level: str = "error", *args, **kwargs):
+        """See the arguments in server_args.py::ServerArgs"""
+
+        # before python program terminates, call shutdown implicitly. Therefore, users don't have to explicitly call .shutdown()
+        atexit.register(self.shutdown)
+
+        server_args = ServerArgs(*args, log_level=log_level, **kwargs)
+        launch_engine(server_args=server_args)
+
+    def generate(
+        self,
+        # The input prompt. It can be a single prompt or a batch of prompts.
+        prompt: Optional[Union[List[str], str]] = None,
+        sampling_params: Optional[Union[List[Dict], Dict]] = None,
+        # The token ids for text; one can either specify text or input_ids.
+        input_ids: Optional[Union[List[List[int]], List[int]]] = None,
+        return_logprob: Optional[Union[List[bool], bool]] = False,
+        logprob_start_len: Optional[Union[List[int], int]] = None,
+        top_logprobs_num: Optional[Union[List[int], int]] = None,
+        lora_path: Optional[List[Optional[str]]] = None,
+        stream: bool = False,
+    ):
+        obj = GenerateReqInput(
+            text=prompt,
+            input_ids=input_ids,
+            sampling_params=sampling_params,
+            return_logprob=return_logprob,
+            logprob_start_len=logprob_start_len,
+            top_logprobs_num=top_logprobs_num,
+            lora_path=lora_path,
+            stream=stream,
+        )
+
+        # get the current event loop
+        loop = asyncio.get_event_loop()
+        ret = loop.run_until_complete(generate_request(obj, None))
+
+        if stream is True:
+
+            def generator_wrapper():
+                offset = 0
+                loop = asyncio.get_event_loop()
+                generator = ret.body_iterator
+                while True:
+                    chunk = loop.run_until_complete(generator.__anext__())
+
+                    if chunk.startswith(STREAM_END_SYMBOL):
+                        break
+                    else:
+                        data = json.loads(chunk[len(STREAM_CHUNK_START_SYMBOL) :])
+                        data["text"] = data["text"][offset:]
+                        offset += len(data["text"])
+                        yield data
+
+            # we cannot yield in the scope of generate() because python does not allow yield + return in the same function
+            # however, it allows to wrap the generator as a subfunction and return
+            return generator_wrapper()
+        else:
+            return ret
+
+    async def async_generate(
+        self,
+        # The input prompt. It can be a single prompt or a batch of prompts.
+        prompt: Optional[Union[List[str], str]] = None,
+        sampling_params: Optional[Dict] = None,
+        # The token ids for text; one can either specify text or input_ids.
+        input_ids: Optional[Union[List[List[int]], List[int]]] = None,
+        return_logprob: Optional[Union[List[bool], bool]] = False,
+        logprob_start_len: Optional[Union[List[int], int]] = None,
+        top_logprobs_num: Optional[Union[List[int], int]] = None,
+        lora_path: Optional[List[Optional[str]]] = None,
+        stream: bool = False,
+    ):
+        obj = GenerateReqInput(
+            text=prompt,
+            input_ids=input_ids,
+            sampling_params=sampling_params,
+            return_logprob=return_logprob,
+            logprob_start_len=logprob_start_len,
+            top_logprobs_num=top_logprobs_num,
+            lora_path=lora_path,
+            stream=stream,
+        )
+
+        ret = await generate_request(obj, None)
+
+        if stream is True:
+            generator = ret.body_iterator
+
+            async def generator_wrapper():
+
+                offset = 0
+
+                while True:
+                    chunk = await generator.__anext__()
+
+                    if chunk.startswith(STREAM_END_SYMBOL):
+                        break
+                    else:
+                        data = json.loads(chunk[len(STREAM_CHUNK_START_SYMBOL) :])
+                        data["text"] = data["text"][offset:]
+                        offset += len(data["text"])
+                        yield data
+
+            return generator_wrapper()
+        else:
+            return ret
+
+    def shutdown(self):
+        kill_process_tree(os.getpid(), include_parent=False)
+
+    def get_tokenizer(self):
+        global tokenizer_manager
+
+        if tokenizer_manager is None:
+            raise ReferenceError("Tokenizer Manager is not initialized.")
+        else:
+            return tokenizer_manager.tokenizer
+
+    def encode(
+        self,
+        prompt: Union[str, List[str], List[Dict], List[List[Dict]]],
+    ):
+        obj = EmbeddingReqInput(text=prompt)
+
+        # get the current event loop
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(encode_request(obj, None))
+
+    def start_profile(self):
+        tokenizer_manager.start_profile()
+
+    def stop_profile(self):
+        tokenizer_manager.stop_profile()
+
+    def get_server_info(self):
+        return {
+            **dataclasses.asdict(tokenizer_manager.server_args),  # server args
+            **scheduler_info,
+            "version": __version__,
+        }
+
+    def init_weights_update_group(
+        self,
+        master_address: str,
+        master_port: int,
+        rank_offset: int,
+        world_size: int,
+        group_name: str,
+        backend: str = "nccl",
+    ):
+        """Initialize parameter update group."""
+        obj = InitWeightsUpdateGroupReqInput(
+            master_address=master_address,
+            master_port=master_port,
+            rank_offset=rank_offset,
+            world_size=world_size,
+            group_name=group_name,
+            backend=backend,
+        )
+
+        async def _init_group():
+            return await tokenizer_manager.init_weights_update_group(obj, None)
+
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(_init_group())
+
+    def update_weights_from_distributed(self, name, dtype, shape):
+        """Update weights from distributed source."""
+        obj = UpdateWeightsFromDistributedReqInput(
+            name=name,
+            dtype=dtype,
+            shape=shape,
+        )
+
+        async def _update_weights():
+            return await tokenizer_manager.update_weights_from_distributed(obj, None)
+
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(_update_weights())
+
+    def get_weights_by_name(self, name, truncate_size=100):
+        """Get weights by parameter name."""
+        obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size)
+
+        async def _get_weights():
+            return await tokenizer_manager.get_weights_by_name(obj, None)
+
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(_get_weights())
+
+
 class Runtime:
     """
-    A wrapper for the server.
+    A wrapper for the HTTP server.
     This is used for launching the server in a python program without
     using the commond line interface.
+
+    It is mainly used for the frontend language.
+    You should use the Engine class if you want to do normal offline processing.
     """
 
     def __init__(
@@ -690,7 +939,7 @@ class Runtime:
 
     def shutdown(self):
         if self.pid is not None:
-            kill_child_process(self.pid, include_self=True)
+            kill_process_tree(self.pid)
             self.pid = None
 
     def cache_prefix(self, prefix: str):
@@ -786,153 +1035,3 @@ class Runtime:
 
     def __del__(self):
         self.shutdown()
-
-
-STREAM_END_SYMBOL = b"data: [DONE]"
-STREAM_CHUNK_START_SYMBOL = b"data:"
-
-
-class Engine:
-    """
-    SRT Engine without an HTTP server layer.
-
-    This class provides a direct inference engine without the need for an HTTP server. It is designed for use cases where
-    launching the HTTP server adds unnecessary complexity or overhead,
-    """
-
-    def __init__(self, *args, **kwargs):
-
-        # before python program terminates, call shutdown implicitly. Therefore, users don't have to explicitly call .shutdown()
-        atexit.register(self.shutdown)
-
-        # runtime server default log level is log
-        # offline engine works in scripts, so we set it to error
-
-        if "log_level" not in kwargs:
-            kwargs["log_level"] = "error"
-
-        server_args = ServerArgs(*args, **kwargs)
-        launch_engine(server_args=server_args)
-
-    def generate(
-        self,
-        # The input prompt. It can be a single prompt or a batch of prompts.
-        prompt: Optional[Union[List[str], str]] = None,
-        sampling_params: Optional[Union[List[Dict], Dict]] = None,
-        # The token ids for text; one can either specify text or input_ids.
-        input_ids: Optional[Union[List[List[int]], List[int]]] = None,
-        return_logprob: Optional[Union[List[bool], bool]] = False,
-        logprob_start_len: Optional[Union[List[int], int]] = None,
-        top_logprobs_num: Optional[Union[List[int], int]] = None,
-        lora_path: Optional[List[Optional[str]]] = None,
-        stream: bool = False,
-    ):
-        obj = GenerateReqInput(
-            text=prompt,
-            input_ids=input_ids,
-            sampling_params=sampling_params,
-            return_logprob=return_logprob,
-            logprob_start_len=logprob_start_len,
-            top_logprobs_num=top_logprobs_num,
-            lora_path=lora_path,
-            stream=stream,
-        )
-
-        # get the current event loop
-        loop = asyncio.get_event_loop()
-        ret = loop.run_until_complete(generate_request(obj, None))
-
-        if stream is True:
-
-            def generator_wrapper():
-                offset = 0
-                loop = asyncio.get_event_loop()
-                generator = ret.body_iterator
-                while True:
-                    chunk = loop.run_until_complete(generator.__anext__())
-
-                    if chunk.startswith(STREAM_END_SYMBOL):
-                        break
-                    else:
-                        data = json.loads(chunk[len(STREAM_CHUNK_START_SYMBOL) :])
-                        data["text"] = data["text"][offset:]
-                        offset += len(data["text"])
-                        yield data
-
-            # we cannot yield in the scope of generate() because python does not allow yield + return in the same function
-            # however, it allows to wrap the generator as a subfunction and return
-            return generator_wrapper()
-        else:
-            return ret
-
-    async def async_generate(
-        self,
-        # The input prompt. It can be a single prompt or a batch of prompts.
-        prompt: Optional[Union[List[str], str]] = None,
-        sampling_params: Optional[Dict] = None,
-        # The token ids for text; one can either specify text or input_ids.
-        input_ids: Optional[Union[List[List[int]], List[int]]] = None,
-        return_logprob: Optional[Union[List[bool], bool]] = False,
-        logprob_start_len: Optional[Union[List[int], int]] = None,
-        top_logprobs_num: Optional[Union[List[int], int]] = None,
-        lora_path: Optional[List[Optional[str]]] = None,
-        stream: bool = False,
-    ):
-        obj = GenerateReqInput(
-            text=prompt,
-            input_ids=input_ids,
-            sampling_params=sampling_params,
-            return_logprob=return_logprob,
-            logprob_start_len=logprob_start_len,
-            top_logprobs_num=top_logprobs_num,
-            lora_path=lora_path,
-            stream=stream,
-        )
-
-        ret = await generate_request(obj, None)
-
-        if stream is True:
-            generator = ret.body_iterator
-
-            async def generator_wrapper():
-
-                offset = 0
-
-                while True:
-                    chunk = await generator.__anext__()
-
-                    if chunk.startswith(STREAM_END_SYMBOL):
-                        break
-                    else:
-                        data = json.loads(chunk[len(STREAM_CHUNK_START_SYMBOL) :])
-                        data["text"] = data["text"][offset:]
-                        offset += len(data["text"])
-                        yield data
-
-            return generator_wrapper()
-        else:
-            return ret
-
-    def shutdown(self):
-        kill_child_process()
-
-    def get_tokenizer(self):
-        global tokenizer_manager
-
-        if tokenizer_manager is None:
-            raise ReferenceError("Tokenizer Manager is not initialized.")
-        else:
-            return tokenizer_manager.tokenizer
-
-    def encode(
-        self,
-        prompt: Union[str, List[str], List[Dict], List[List[Dict]]],
-    ):
-        obj = EmbeddingReqInput(text=prompt)
-
-        # get the current event loop
-        loop = asyncio.get_event_loop()
-        return loop.run_until_complete(encode_request(obj, None))
-
-    async def get_server_info(self):
-        return await _get_server_info()