Allow multi SGLang engines to coordinate

examples/runtime/engine/offline_batch_inference_torchrun.py

@@ -0,0 +1,79 @@
+import datetime
+import os
+import sys
+
+from sglang.srt.server.engine_fragment import EngineFragment
+
+
+def run():
+    """
+    Example command:
+    ```
+    torchrun --nproc_per_node=4 offline_batch_inference_torchrun.py
+    ```
+    """
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    rank = int(os.environ["RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+
+    def _log(text):
+        t = datetime.datetime.now().strftime("%H:%M:%S")
+        print(f"[{t}] [rank={rank}] {text}")
+
+    _log(
+        f'start {local_rank=} {rank=} {world_size=} {sys.argv=} {os.environ.get("CUDA_VISIBLE_DEVICES")}'
+    )
+
+    tp_size = world_size
+    tp_rank = rank
+    _log(f"{tp_rank=} {tp_size=}")
+
+    model_name, mem_fraction_static = "meta-llama/Llama-3.2-1B-Instruct", 0.1
+    # model_name, mem_fraction_static = "meta-llama/Llama-3.1-70B-Instruct", 0.9 # test large models
+
+    for k in [
+        "GROUP_RANK",
+        "GROUP_WORLD_SIZE",
+        "LOCAL_RANK",
+        "LOCAL_WORLD_SIZE",
+        "MASTER_ADDR",
+        "MASTER_PORT",
+        "NCCL_DEBUG",
+        "OMP_NUM_THREADS",
+        "RANK",
+        "ROLE_NAME",
+        "ROLE_RANK",
+        "ROLE_WORLD_SIZE",
+        "TORCHELASTIC_ERROR_FILE",
+        "TORCHELASTIC_MAX_RESTARTS",
+        "TORCHELASTIC_RESTART_COUNT",
+        "TORCHELASTIC_RUN_ID",
+        "TORCHELASTIC_USE_AGENT_STORE",
+        "TORCH_NCCL_ASYNC_ERROR_HANDLING",
+        "WORLD_SIZE",
+    ]:
+        del os.environ[k]
+
+    fragment = EngineFragment(
+        model_path=model_name,
+        mem_fraction_static=mem_fraction_static,
+        tp_size=tp_size,
+        tp_rank=tp_rank,
+        nccl_port=23456,
+        gpu_id=tp_rank,
+    )
+    _log(f"{fragment=}")
+
+    output = fragment.generate(
+        prompt=["1+1=2, 1+2=3, 1+3=4, 1+4=", "9-1=8, 8-1=7, 7-1="],
+        sampling_params=dict(max_new_tokens=16, temperature=0.0),
+    )
+    _log(f"{output=}")
+
+    fragment.shutdown()
+    _log(f"End script")
+
+
+if __name__ == "__main__":
+    run()

python/sglang/bench_one_batch.py

@@ -60,8 +60,8 @@
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.orchestration.std.launcher import _set_envs_and_config
 from sglang.srt.sampling.sampling_params import SamplingParams
-from sglang.srt.server import _set_envs_and_config
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import configure_logger, kill_process_tree, suppress_other_loggers

python/sglang/launch_server_llavavid.py

@@ -3,7 +3,8 @@
 import json
 import sys
 
-from sglang.srt.server import launch_server, prepare_server_args
+from sglang.srt.server import launch_server
+from sglang.srt.server_args import prepare_server_args
 
 if __name__ == "__main__":
     server_args = prepare_server_args(sys.argv[1:])

python/sglang/srt/managers/data_parallel_controller.py

@@ -27,7 +27,7 @@
     TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
 )
-from sglang.srt.managers.scheduler import run_scheduler_process
+from sglang.srt.orchestration.std.scheduler import run_scheduler_process
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import bind_port, configure_logger, get_zmq_socket
 from sglang.utils import get_exception_traceback

python/sglang/srt/managers/detokenizer_manager.py

@@ -15,23 +15,17 @@
 
 import dataclasses
 import logging
-import signal
 from collections import OrderedDict
 from typing import Dict, List, Union
 
-import psutil
-import setproctitle
-import zmq
-
 from sglang.srt.hf_transformers_utils import get_tokenizer
 from sglang.srt.managers.io_struct import (
     BatchEmbeddingOut,
     BatchStrOut,
     BatchTokenIDOut,
 )
-from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import configure_logger, get_zmq_socket
-from sglang.utils import find_printable_text, get_exception_traceback
+from sglang.srt.server_args import ServerArgs
+from sglang.utils import find_printable_text
 
 logger = logging.getLogger(__name__)
 
@@ -53,17 +47,7 @@ class DetokenizerManager:
     def __init__(
         self,
         server_args: ServerArgs,
-        port_args: PortArgs,
     ):
-        # Init inter-process communication
-        context = zmq.Context(2)
-        self.recv_from_scheduler = get_zmq_socket(
-            context, zmq.PULL, port_args.detokenizer_ipc_name
-        )
-        self.send_to_tokenizer = get_zmq_socket(
-            context, zmq.PUSH, port_args.tokenizer_ipc_name
-        )
-
         if server_args.skip_tokenizer_init:
             self.tokenizer = None
         else:
@@ -75,125 +59,116 @@ def __init__(
 
         self.decode_status = LimitedCapacityDict()
 
-    def trim_matched_stop(
-        self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool
-    ):
-        if no_stop_trim or not finished_reason:
-            return output
+    def handle_batch_embedding_out(self, recv_obj: BatchEmbeddingOut):
+        # If it is embedding model, no detokenization is needed.
+        return recv_obj
+
+    def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOut):
+        bs = len(recv_obj.rids)
+
+        # Initialize decode status
+        read_ids, surr_ids = [], []
+        for i in range(bs):
+            rid = recv_obj.rids[i]
+            vid = recv_obj.vids[i]
+            if rid not in self.decode_status or self.decode_status[rid].vid != vid:
+                s = DecodeStatus(
+                    vid=vid,
+                    decoded_text=recv_obj.decoded_texts[i],
+                    decode_ids=recv_obj.decode_ids[i],
+                    surr_offset=0,
+                    read_offset=recv_obj.read_offsets[i],
+                )
+                self.decode_status[rid] = s
+            else:
+                s = self.decode_status[rid]
+                s.decode_ids = recv_obj.decode_ids[i]
+
+            read_ids.append(
+                _trim_matched_stop(
+                    s.decode_ids[s.surr_offset :],
+                    recv_obj.finished_reasons[i],
+                    recv_obj.no_stop_trim[i],
+                )
+            )
+            surr_ids.append(s.decode_ids[s.surr_offset : s.read_offset])
 
-        matched = finished_reason.get("matched", None)
-        if not matched:
-            return output
+        # TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request
+        surr_texts = self.tokenizer.batch_decode(
+            surr_ids,
+            skip_special_tokens=recv_obj.skip_special_tokens[0],
+            spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
+        )
+        read_texts = self.tokenizer.batch_decode(
+            read_ids,
+            skip_special_tokens=recv_obj.skip_special_tokens[0],
+            spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
+        )
 
-        # TODO(lmzheng): handle the case where multiple stop strs are hit
+        # Incremental decoding
+        output_strs = []
+        for i in range(bs):
+            s = self.decode_status[recv_obj.rids[i]]
+            new_text = read_texts[i][len(surr_texts[i]) :]
+            if recv_obj.finished_reasons[i] is None:
+                # Streaming chunk: update the decode status
+                if len(new_text) > 0 and not new_text.endswith("�"):
+                    s.decoded_text = s.decoded_text + new_text
+                    s.surr_offset = s.read_offset
+                    s.read_offset = len(s.decode_ids)
+                    new_text = ""
+                else:
+                    new_text = find_printable_text(new_text)
 
-        # Trim stop str.
-        if isinstance(matched, str) and isinstance(output, str):
-            pos = output.find(matched)
-            return output[:pos] if pos != -1 else output
+            output_strs.append(
+                _trim_matched_stop(
+                    s.decoded_text + new_text,
+                    recv_obj.finished_reasons[i],
+                    recv_obj.no_stop_trim[i],
+                )
+            )
 
-        # Trim stop token.
-        if isinstance(matched, int) and isinstance(output, list):
-            assert len(output) > 0
-            return output[:-1]
-        return output
+        return BatchStrOut(
+            rids=recv_obj.rids,
+            finished_reasons=recv_obj.finished_reasons,
+            output_strs=output_strs,
+            prompt_tokens=recv_obj.prompt_tokens,
+            completion_tokens=recv_obj.completion_tokens,
+            cached_tokens=recv_obj.cached_tokens,
+            input_token_logprobs_val=recv_obj.input_token_logprobs_val,
+            input_token_logprobs_idx=recv_obj.input_token_logprobs_idx,
+            output_token_logprobs_val=recv_obj.output_token_logprobs_val,
+            output_token_logprobs_idx=recv_obj.output_token_logprobs_idx,
+            input_top_logprobs_val=recv_obj.input_top_logprobs_val,
+            input_top_logprobs_idx=recv_obj.input_top_logprobs_idx,
+            output_top_logprobs_val=recv_obj.output_top_logprobs_val,
+            output_top_logprobs_idx=recv_obj.output_top_logprobs_idx,
+            normalized_prompt_logprob=recv_obj.normalized_prompt_logprob,
+        )
 
-    def event_loop(self):
-        """The event loop that handles requests"""
 
-        while True:
-            recv_obj = self.recv_from_scheduler.recv_pyobj()
+def _trim_matched_stop(
+    output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool
+):
+    if no_stop_trim or not finished_reason:
+        return output
 
-            if isinstance(recv_obj, BatchEmbeddingOut):
-                # If it is embedding model, no detokenization is needed.
-                self.send_to_tokenizer.send_pyobj(recv_obj)
-                continue
-            else:
-                assert isinstance(recv_obj, BatchTokenIDOut)
-
-            bs = len(recv_obj.rids)
-
-            # Initialize decode status
-            read_ids, surr_ids = [], []
-            for i in range(bs):
-                rid = recv_obj.rids[i]
-                vid = recv_obj.vids[i]
-                if rid not in self.decode_status or self.decode_status[rid].vid != vid:
-                    s = DecodeStatus(
-                        vid=vid,
-                        decoded_text=recv_obj.decoded_texts[i],
-                        decode_ids=recv_obj.decode_ids[i],
-                        surr_offset=0,
-                        read_offset=recv_obj.read_offsets[i],
-                    )
-                    self.decode_status[rid] = s
-                else:
-                    s = self.decode_status[rid]
-                    s.decode_ids = recv_obj.decode_ids[i]
-
-                read_ids.append(
-                    self.trim_matched_stop(
-                        s.decode_ids[s.surr_offset :],
-                        recv_obj.finished_reasons[i],
-                        recv_obj.no_stop_trim[i],
-                    )
-                )
-                surr_ids.append(s.decode_ids[s.surr_offset : s.read_offset])
+    matched = finished_reason.get("matched", None)
+    if not matched:
+        return output
 
-            # TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request
-            surr_texts = self.tokenizer.batch_decode(
-                surr_ids,
-                skip_special_tokens=recv_obj.skip_special_tokens[0],
-                spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
-            )
-            read_texts = self.tokenizer.batch_decode(
-                read_ids,
-                skip_special_tokens=recv_obj.skip_special_tokens[0],
-                spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
-            )
+    # TODO(lmzheng): handle the case where multiple stop strs are hit
 
-            # Incremental decoding
-            output_strs = []
-            for i in range(bs):
-                s = self.decode_status[recv_obj.rids[i]]
-                new_text = read_texts[i][len(surr_texts[i]) :]
-                if recv_obj.finished_reasons[i] is None:
-                    # Streaming chunk: update the decode status
-                    if len(new_text) > 0 and not new_text.endswith("�"):
-                        s.decoded_text = s.decoded_text + new_text
-                        s.surr_offset = s.read_offset
-                        s.read_offset = len(s.decode_ids)
-                        new_text = ""
-                    else:
-                        new_text = find_printable_text(new_text)
-
-                output_strs.append(
-                    self.trim_matched_stop(
-                        s.decoded_text + new_text,
-                        recv_obj.finished_reasons[i],
-                        recv_obj.no_stop_trim[i],
-                    )
-                )
+    # Trim stop str.
+    if isinstance(matched, str) and isinstance(output, str):
+        pos = output.find(matched)
+        return output[:pos] if pos != -1 else output
 
-            self.send_to_tokenizer.send_pyobj(
-                BatchStrOut(
-                    rids=recv_obj.rids,
-                    finished_reasons=recv_obj.finished_reasons,
-                    output_strs=output_strs,
-                    prompt_tokens=recv_obj.prompt_tokens,
-                    completion_tokens=recv_obj.completion_tokens,
-                    cached_tokens=recv_obj.cached_tokens,
-                    input_token_logprobs_val=recv_obj.input_token_logprobs_val,
-                    input_token_logprobs_idx=recv_obj.input_token_logprobs_idx,
-                    output_token_logprobs_val=recv_obj.output_token_logprobs_val,
-                    output_token_logprobs_idx=recv_obj.output_token_logprobs_idx,
-                    input_top_logprobs_val=recv_obj.input_top_logprobs_val,
-                    input_top_logprobs_idx=recv_obj.input_top_logprobs_idx,
-                    output_top_logprobs_val=recv_obj.output_top_logprobs_val,
-                    output_top_logprobs_idx=recv_obj.output_top_logprobs_idx,
-                    normalized_prompt_logprob=recv_obj.normalized_prompt_logprob,
-                )
-            )
+    # Trim stop token.
+    if isinstance(matched, int) and isinstance(output, list):
+        assert len(output) > 0
+        return output[:-1]
+    return output
 
 
 class LimitedCapacityDict(OrderedDict):
@@ -207,20 +182,3 @@ def __setitem__(self, key, value):
             self.popitem(last=False)
         # Set the new item
         super().__setitem__(key, value)
-
-
-def run_detokenizer_process(
-    server_args: ServerArgs,
-    port_args: PortArgs,
-):
-    setproctitle.setproctitle("sglang::detokenizer")
-    configure_logger(server_args)
-    parent_process = psutil.Process().parent()
-
-    try:
-        manager = DetokenizerManager(server_args, port_args)
-        manager.event_loop()
-    except Exception:
-        traceback = get_exception_traceback()
-        logger.error(f"DetokenizerManager hit an exception: {traceback}")
-        parent_process.send_signal(signal.SIGQUIT)