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

sglang/srt/utils.py

@@ -1,24 +1,31 @@
 """Common utilities."""
 
 import base64
+import fcntl
+import logging
+import multiprocessing
 import os
 import random
 import socket
+import struct
 import time
 from importlib.metadata import PackageNotFoundError, version
 from io import BytesIO
 from typing import List, Optional
 
 import numpy as np
-import pydantic
+import psutil
 import requests
+import rpyc
 import torch
+import triton
 from fastapi.responses import JSONResponse
 from packaging import version as pkg_version
-from pydantic import BaseModel
+from rpyc.utils.server import ThreadedServer
 from starlette.middleware.base import BaseHTTPMiddleware
 
-from sglang.utils import get_exception_traceback
+logger = logging.getLogger(__name__)
+
 
 show_time_cost = False
 time_infos = {}
@@ -90,7 +97,7 @@ def calculate_time(show=False, min_cost_ms=0.0):
     return wrapper
 
 
-def get_available_gpu_memory(gpu_id, distributed=True):
+def get_available_gpu_memory(gpu_id, distributed=False):
     """
     Get available memory for cuda:gpu_id device.
     When distributed is True, the available memory is the minimum available memory of all GPUs.
@@ -104,6 +111,7 @@ def get_available_gpu_memory(gpu_id, distributed=True):
             "which may cause useless memory allocation for torch CUDA context.",
         )
 
+    torch.cuda.empty_cache()
     free_gpu_memory, _ = torch.cuda.mem_get_info(gpu_id)
 
     if distributed:
@@ -117,38 +125,21 @@ def get_available_gpu_memory(gpu_id, distributed=True):
 
 
 def set_random_seed(seed: int) -> None:
+    """Set the random seed for all libraries."""
     random.seed(seed)
-
+    np.random.seed(seed)
     torch.manual_seed(seed)
     if torch.cuda.is_available():
         torch.cuda.manual_seed_all(seed)
 
 
-def alloc_usable_network_port(num, used_list=()):
-    port_list = []
-    for port in range(10000, 65536):
-        if port in used_list:
-            continue
-
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-            try:
-                s.bind(("", port))
-                s.listen(1)  # Attempt to listen on the port
-                port_list.append(port)
-            except socket.error:
-                pass  # If any error occurs, this port is not usable
-
-            if len(port_list) == num:
-                return port_list
-    return None
-
-
-def check_port(port):
+def is_port_available(port):
+    """Return whether a port is available."""
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
         try:
             s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
             s.bind(("", port))
+            s.listen(1)
             return True
         except socket.error:
             return False
@@ -158,35 +149,34 @@ def allocate_init_ports(
     port: Optional[int] = None,
     additional_ports: Optional[List[int]] = None,
     tp_size: int = 1,
+    dp_size: int = 1,
 ):
-    port = 30000 if port is None else port
-    additional_ports = [] if additional_ports is None else additional_ports
-    additional_ports = (
-        [additional_ports] if isinstance(additional_ports, int) else additional_ports
-    )
-    # first check on server port
-    if not check_port(port):
-        new_port = alloc_usable_network_port(1, used_list=[port])[0]
-        print(f"WARNING: Port {port} is not available. Use {new_port} instead.")
-        port = new_port
-
-    # then we check on additional ports
-    additional_unique_ports = set(additional_ports) - {port}
-    # filter out ports that are already in use
-    can_use_ports = [port for port in additional_unique_ports if check_port(port)]
-
-    num_specified_ports = len(can_use_ports)
-    if num_specified_ports < 4 + tp_size:
-        addtional_can_use_ports = alloc_usable_network_port(
-            num=4 + tp_size - num_specified_ports, used_list=can_use_ports + [port]
+    """Allocate ports for all connections."""
+    if additional_ports:
+        ret_ports = [port] + additional_ports
+    else:
+        ret_ports = [port]
+
+    ret_ports = list(set(x for x in ret_ports if is_port_available(x)))
+    cur_port = ret_ports[-1] + 1 if len(ret_ports) > 0 else 10000
+
+    # HTTP + Tokenizer + Controller + Detokenizer + dp_size * (nccl + tp_size)
+    num_ports_needed = 4 + dp_size * (1 + tp_size)
+    while len(ret_ports) < num_ports_needed:
+        if cur_port not in ret_ports and is_port_available(cur_port):
+            ret_ports.append(cur_port)
+        cur_port += 1
+
+    if port is not None and ret_ports[0] != port:
+        logger.warn(
+            f"WARNING: Port {port} is not available. Use port {ret_ports[0]} instead."
         )
-        can_use_ports.extend(addtional_can_use_ports)
 
-    additional_ports = can_use_ports[: 4 + tp_size]
-    return port, additional_ports
+    return ret_ports[0], ret_ports[1:num_ports_needed]
 
 
 def get_int_token_logit_bias(tokenizer, vocab_size):
+    """Get the logit bias for integer-only tokens."""
     # a bug when model's vocab size > tokenizer.vocab_size
     vocab_size = tokenizer.vocab_size
     logit_bias = np.zeros(vocab_size, dtype=np.float32)
@@ -200,14 +190,11 @@ def get_int_token_logit_bias(tokenizer, vocab_size):
 
 def wrap_kernel_launcher(kernel):
     """A faster launcher for triton kernels."""
-    import torch.distributed as dist
-
-    if dist.is_initialized():
-        rank = dist.get_rank()
-    else:
-        rank = 0
+    if int(triton.__version__.split(".")[0]) >= 3:
+        return None
 
-    kernels = kernel.cache[rank].values()
+    gpu_id = torch.cuda.current_device()
+    kernels = kernel.cache[gpu_id].values()
     kernel = next(iter(kernels))
 
     # Different trition versions use different low-level names
@@ -275,7 +262,9 @@ def is_multimodal_model(model):
 
     if isinstance(model, ModelConfig):
         model_path = model.path.lower()
-        return "llava" in model_path or "yi-vl" in model_path or "llava-next" in model_path
+        return (
+            "llava" in model_path or "yi-vl" in model_path or "llava-next" in model_path
+        )
 
     raise ValueError("unrecognized type")
 
@@ -382,20 +371,145 @@ def load_image(image_file):
     return image, image_size
 
 
-def assert_pkg_version(pkg: str, min_version: str):
+def connect_rpyc_service(host, port):
+    repeat_count = 0
+    while repeat_count < 20:
+        try:
+            con = rpyc.connect(
+                host,
+                port,
+                config={
+                    "allow_public_attrs": True,
+                    "allow_pickle": True,
+                    "sync_request_timeout": 3600,
+                },
+            )
+            break
+        except ConnectionRefusedError as e:
+            time.sleep(1)
+        repeat_count += 1
+    if repeat_count == 20:
+        raise RuntimeError(f"Connect rpyc error: {e}")
+
+    return con.root
+
+
+def start_rpyc_service(service: rpyc.Service, port: int):
+    t = ThreadedServer(
+        service=service,
+        port=port,
+        protocol_config={
+            "allow_public_attrs": True,
+            "allow_pickle": True,
+            "sync_request_timeout": 3600,
+        },
+    )
+    t.logger.setLevel(logging.WARN)
+    t.start()
+
+
+def start_rpyc_service_process(service: rpyc.Service, port: int):
+    proc = multiprocessing.Process(target=start_rpyc_service, args=(service, port))
+    proc.start()
+    return proc
+
+
+def suppress_other_loggers():
+    from vllm.logger import logger as vllm_default_logger
+
+    vllm_default_logger.setLevel(logging.WARN)
+    logging.getLogger("vllm.config").setLevel(logging.ERROR)
+    logging.getLogger("vllm.distributed.device_communicators.pynccl").setLevel(
+        logging.WARN
+    )
+    logging.getLogger("vllm.selector").setLevel(logging.WARN)
+    logging.getLogger("vllm.utils").setLevel(logging.WARN)
+
+
+def assert_pkg_version(pkg: str, min_version: str, message: str):
     try:
         installed_version = version(pkg)
         if pkg_version.parse(installed_version) < pkg_version.parse(min_version):
             raise Exception(
-                f"{pkg} is installed with version {installed_version} which "
-                f"is less than the minimum required version {min_version}"
+                f"{pkg} is installed with version {installed_version}, which "
+                f"is less than the minimum required version {min_version}. " +
+                message
             )
     except PackageNotFoundError:
         raise Exception(
-            f"{pkg} with minimum required version {min_version} is not installed"
+            f"{pkg} with minimum required version {min_version} is not installed. " +
+            message
         )
 
 
+def kill_parent_process():
+    """Kill the parent process and all children of the parent process."""
+    current_process = psutil.Process()
+    parent_process = current_process.parent()
+    children = current_process.children(recursive=True)
+    for child in children:
+        if child.pid != current_process.pid:
+            os.kill(child.pid, 9)
+    os.kill(parent_process.pid, 9)
+
+
+def monkey_patch_vllm_p2p_access_check(gpu_id: int):
+    """
+    Monkey patch the slow p2p access check in vllm.
+    NOTE: We assume the p2p access is always allowed, which can be wrong for some setups.
+    """
+
+    # TODO: need a better check than just dev str name match
+    # compat: skip RTX 40 series as they do not have P2P feature and even checking for them may cause errors
+    device_name = torch.cuda.get_device_name(gpu_id)
+    if "RTX 40" not in device_name:
+        import vllm.distributed.device_communicators.custom_all_reduce_utils as tgt
+
+        setattr(tgt, "gpu_p2p_access_check", lambda *arg, **kwargs: True)
+
+
+def monkey_patch_vllm_dummy_weight_loader():
+    """
+    Monkey patch the dummy weight loader in vllm to call process_weights_after_loading.
+    """
+
+    from vllm.model_executor.model_loader.loader import (
+        ModelConfig, DeviceConfig, LoRAConfig, VisionLanguageConfig,
+        ParallelConfig, SchedulerConfig, CacheConfig, nn,
+        set_default_torch_dtype, _initialize_model, initialize_dummy_weights,
+        DummyModelLoader
+    )
+
+    def load_model(self, *, model_config: ModelConfig,
+                   device_config: DeviceConfig,
+                   lora_config: Optional[LoRAConfig],
+                   vision_language_config: Optional[VisionLanguageConfig],
+                   parallel_config: ParallelConfig,
+                   scheduler_config: SchedulerConfig,
+                   cache_config: CacheConfig) -> nn.Module:
+        with set_default_torch_dtype(model_config.dtype):
+            with torch.device(device_config.device):
+                model = _initialize_model(model_config, self.load_config,
+                                          lora_config, vision_language_config,
+                                          cache_config)
+
+            for _, module in model.named_modules():
+                quant_method = getattr(module, "quant_method", None)
+                if quant_method is not None:
+                    quant_method.process_weights_after_loading(module)
+                # FIXME: Remove this after Mixtral is updated
+                # to use quant_method.
+                if hasattr(module, "process_weights_after_loading"):
+                    module.process_weights_after_loading()
+
+            # NOTE(woosuk): For accurate performance evaluation, we assign
+            # random values to the weights.
+            initialize_dummy_weights(model)
+        return model.eval()
+
+    setattr(DummyModelLoader, "load_model", load_model)
+
+
 API_KEY_HEADER_NAME = "X-API-Key"
 
 
@@ -416,11 +530,64 @@ class APIKeyValidatorMiddleware(BaseHTTPMiddleware):
         return response
 
 
-# FIXME: Remove this once we drop support for pydantic 1.x
-IS_PYDANTIC_1 = int(pydantic.VERSION.split(".")[0]) == 1
+def get_ip_address(ifname):
+    """
+    Get the IP address of a network interface.
+
+    :param ifname: Name of the network interface (e.g., 'eth0')
+    :return: IP address of the network interface
+    """
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    ip_address = fcntl.ioctl(
+        s.fileno(),
+        0x8915,  # SIOCGIFADDR
+        struct.pack('256s', bytes(ifname[:15], 'utf-8'))
+    )[20:24]
+    return socket.inet_ntoa(ip_address)
+
+
+def send_addrs_to_rank_0(model_port_args, server_args):
+    assert server_args.node_rank != 0 and server_args.dp_size == 1
+    import torch.distributed as dist
+
+    ifname = os.environ.get("SGLANG_SOCKET_IFNAME", os.environ.get("NCCL_SOCKET_IFNAME", "eth0"))
+    ip_addr = get_ip_address(ifname)
+
+    num_tp_ports = server_args.tp_size // server_args.nnodes
+    model_port_args.model_tp_ips[:num_tp_ports] = [ip_addr] * num_tp_ports
+    ip_addr = [int(x) for x in ip_addr.split(".")]
+    addrs_tensor = torch.tensor(ip_addr + model_port_args.model_tp_ports, dtype=torch.int)
+
+    init_method = f"tcp://{server_args.nccl_init_addr}"
+    dist.init_process_group(backend="gloo", init_method=init_method, rank=server_args.node_rank, world_size=server_args.nnodes)
+    dist.send(addrs_tensor, dst=0)
+    print(f"Node {server_args.node_rank} sent: ip_address {ip_addr} and ports {model_port_args.model_tp_ports}")
+
+    dist.barrier()
+    dist.destroy_process_group() 
+
+
+def receive_addrs(model_port_args, server_args):
+    assert server_args.node_rank == 0 and server_args.dp_size == 1
+    import torch.distributed as dist
+
+    ifname = os.environ.get("SGLANG_SOCKET_IFNAME", os.environ.get("NCCL_SOCKET_IFNAME", "eth0"))
+    ip_addr = get_ip_address(ifname)
+
+    num_tp_ports = server_args.tp_size // server_args.nnodes
+    model_port_args.model_tp_ips[:num_tp_ports] = [ip_addr] * num_tp_ports
+
+    init_method = f"tcp://{server_args.nccl_init_addr}"
+    dist.init_process_group(backend="gloo", init_method=init_method, rank=server_args.node_rank, world_size=server_args.nnodes)
 
+    for src_rank in range(1, server_args.nnodes):
+        tensor = torch.zeros(4 + num_tp_ports, dtype=torch.int)
+        dist.recv(tensor, src=src_rank)
+        ip = ".".join([str(x) for x in tensor[:4].tolist()])
+        ports = tensor[4:].tolist()
+        model_port_args.model_tp_ips[num_tp_ports * src_rank: num_tp_ports * (src_rank + 1)] = [ip] * num_tp_ports
+        model_port_args.model_tp_ports[num_tp_ports * src_rank: num_tp_ports * (src_rank + 1)] = ports
+        print(f"Node 0 received from rank {src_rank}: {tensor.tolist()}")
 
-def jsonify_pydantic_model(obj: BaseModel):
-    if IS_PYDANTIC_1:
-        return obj.json(ensure_ascii=False)
-    return obj.model_dump_json()
+    dist.barrier()
+    dist.destroy_process_group() 

sglang/test/test_programs.py

@@ -1,6 +1,4 @@
-"""
-This file contains the SGL programs used for unit testing.
-"""
+"""This file contains the SGL programs used for unit testing."""
 
 import json
 import re
@@ -304,6 +302,7 @@ def test_image_qa():
         temperature=0,
         max_new_tokens=64,
     )
+
     assert (
         "taxi" in state.messages()[-1]["content"]
         or "car" in state.messages()[-1]["content"]
@@ -349,3 +348,66 @@ def test_regex():
     state = regex_gen.run()
     answer = state["answer"]
     assert re.match(regex, answer)
+
+
+def test_completion_speculative():
+    @sgl.function(num_api_spec_tokens=64)
+    def gen_character_spec(s):
+        s += "Construct a character within the following format:\n"
+        s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n"
+        s += "\nPlease generate new Name, Birthday and Job.\n"
+        s += (
+            "Name:"
+            + sgl.gen("name", stop="\n")
+            + "\nBirthday:"
+            + sgl.gen("birthday", stop="\n")
+        )
+        s += "\nJob:" + sgl.gen("job", stop="\n") + "\n"
+
+    @sgl.function
+    def gen_character_no_spec(s):
+        s += "Construct a character within the following format:\n"
+        s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n"
+        s += "\nPlease generate new Name, Birthday and Job.\n"
+        s += (
+            "Name:"
+            + sgl.gen("name", stop="\n")
+            + "\nBirthday:"
+            + sgl.gen("birthday", stop="\n")
+        )
+        s += "\nJob:" + sgl.gen("job", stop="\n") + "\n"
+
+    token_usage = sgl.global_config.default_backend.token_usage
+
+    token_usage.reset()
+    gen_character_spec().sync()
+    usage_with_spec = token_usage.prompt_tokens
+
+    token_usage.reset()
+    gen_character_no_spec().sync()
+    usage_with_no_spec = token_usage.prompt_tokens
+
+    assert (
+        usage_with_spec < usage_with_no_spec
+    ), f"{usage_with_spec} vs {usage_with_no_spec}"
+
+
+def test_chat_completion_speculative():
+    @sgl.function(num_api_spec_tokens=256)
+    def gen_character_spec(s):
+        s += sgl.system("You are a helpful assistant.")
+        s += sgl.user("Construct a character within the following format:")
+        s += sgl.assistant(
+            "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n"
+        )
+        s += sgl.user("Please generate new Name, Birthday and Job.\n")
+        s += sgl.assistant(
+            "Name:"
+            + sgl.gen("name", stop="\n")
+            + "\nBirthday:"
+            + sgl.gen("birthday", stop="\n")
+            + "\nJob:"
+            + sgl.gen("job", stop="\n")
+        )
+
+    gen_character_spec().sync()

sglang/test/test_utils.py

@@ -9,7 +9,7 @@ import requests
 from sglang.backend.openai import OpenAI
 from sglang.backend.runtime_endpoint import RuntimeEndpoint
 from sglang.global_config import global_config
-from sglang.srt.utils import get_exception_traceback
+from sglang.utils import get_exception_traceback
 
 
 def call_generate_lightllm(prompt, temperature, max_tokens, stop=None, url=None):
@@ -88,6 +88,33 @@ def call_generate_srt_raw(prompt, temperature, max_tokens, stop=None, url=None):
     return pred
 
 
+def call_generate_ginfer(prompt, temperature, max_tokens, stop=None, url=None):
+    import grpc
+    from ginfer import sampler_pb2, sampler_pb2_grpc
+
+    sampler_channel = grpc.insecure_channel(url.replace("http://", ""))
+    sampler = sampler_pb2_grpc.SamplerStub(sampler_channel)
+
+    if stop is None:
+        stop_strings = None
+    else:
+        stop_strings = [stop]
+
+    sample_request = sampler_pb2.SampleTextRequest(
+        prompt=prompt,
+        settings=sampler_pb2.SampleSettings(
+            max_len=max_tokens,
+            rng_seed=0,
+            temperature=max(temperature, 1e-7),
+            nucleus_p=1,
+            stop_strings=stop_strings,
+        ),
+    )
+    stream = sampler.SampleText(sample_request)
+    response = "".join([x.text for x in stream])
+    return response
+
+
 def call_generate_guidance(
     prompt, temperature, max_tokens, stop=None, n=1, regex=None, model=None
 ):
@@ -228,6 +255,7 @@ def add_common_other_args_and_parse(parser):
             "vllm",
             "outlines",
             "lightllm",
+            "ginfer",
             "guidance",
             "lmql",
             "srt-raw",
@@ -248,6 +276,7 @@ def add_common_other_args_and_parse(parser):
             "lightllm": 22000,
             "lmql": 23000,
             "srt-raw": 30000,
+            "ginfer": 9988,
         }
         args.port = default_port.get(args.backend, None)
     return args
@@ -283,6 +312,8 @@ def _get_call_generate(args):
         return partial(call_generate_vllm, url=f"{args.host}:{args.port}/generate")
     elif args.backend == "srt-raw":
         return partial(call_generate_srt_raw, url=f"{args.host}:{args.port}/generate")
+    elif args.backend == "ginfer":
+        return partial(call_generate_ginfer, url=f"{args.host}:{args.port}")
     elif args.backend == "outlines":
         return partial(call_generate_outlines, url=f"{args.host}:{args.port}/generate")
     elif args.backend == "guidance":

sglang/utils.py

@@ -2,7 +2,8 @@
 
 import base64
 import json
-import os
+import logging
+import signal
 import sys
 import threading
 import traceback
@@ -14,6 +15,8 @@ from json import dumps
 import numpy as np
 import requests
 
+logger = logging.getLogger(__name__)
+
 
 def get_exception_traceback():
     etype, value, tb = sys.exc_info()
@@ -93,8 +96,12 @@ def http_request(
             data = None
         else:
             data = bytes(dumps(json), encoding="utf-8")
-        resp = urllib.request.urlopen(req, data=data, cafile=verify)
-        return HttpResponse(resp)
+
+        try:
+            resp = urllib.request.urlopen(req, data=data, cafile=verify)
+            return HttpResponse(resp)
+        except urllib.error.HTTPError as e:
+            return HttpResponse(e)
 
 
 def encode_image_base64(image_path):
@@ -137,7 +144,8 @@ def encode_frame(frame):
 
 
 def encode_video_base64(video_path, num_frames=16):
-    import cv2
+    import cv2  # pip install opencv-python-headless
+
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         raise IOError(f"Could not open video file:{video_path}")
@@ -242,3 +250,14 @@ def run_with_timeout(func, args=(), kwargs=None, timeout=None):
         raise RuntimeError()
 
     return ret_value[0]
+
+
+def graceful_registry(sub_module_name):
+    def graceful_shutdown(signum, frame):
+        logger.info(
+            f"{sub_module_name} Received signal to shutdown. Performing graceful shutdown..."
+        )
+        if signum == signal.SIGTERM:
+            logger.info(f"{sub_module_name} recive sigterm")
+
+    signal.signal(signal.SIGTERM, graceful_shutdown)