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

sglang/srt/utils.py

@@ -1,25 +1,31 @@
 """Common utilities."""
 
 import base64
+import multiprocessing
+import logging
 import os
 import random
 import socket
-import sys
 import time
-import traceback
 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 rpyc.utils.server import ThreadedServer
 from fastapi.responses import JSONResponse
 from packaging import version as pkg_version
-from pydantic import BaseModel
 from starlette.middleware.base import BaseHTTPMiddleware
 
+
+logger = logging.getLogger(__name__)
+
+
 show_time_cost = False
 time_infos = {}
 
@@ -90,37 +96,49 @@ def calculate_time(show=False, min_cost_ms=0.0):
     return wrapper
 
 
-def set_random_seed(seed: int) -> None:
-    random.seed(seed)
+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.
+    """
+    num_gpus = torch.cuda.device_count()
+    assert gpu_id < num_gpus
 
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
+    if torch.cuda.current_device() != gpu_id:
+        print(
+            f"WARNING: current device is not {gpu_id}, but {torch.cuda.current_device()}, ",
+            "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:
+        tensor = torch.tensor(free_gpu_memory, dtype=torch.float32).to(
+            torch.device("cuda", gpu_id)
+        )
+        torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.MIN)
+        free_gpu_memory = tensor.item()
 
-def alloc_usable_network_port(num, used_list=()):
-    port_list = []
-    for port in range(10000, 65536):
-        if port in used_list:
-            continue
+    return free_gpu_memory / (1 << 30)
 
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            try:
-                s.bind(("", port))
-                port_list.append(port)
-            except socket.error:
-                pass
 
-            if len(port_list) == num:
-                return port_list
-    return None
+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 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
@@ -130,41 +148,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]
-        )
-        can_use_ports.extend(addtional_can_use_ports)
+    """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
 
-    additional_ports = can_use_ports[: 4 + tp_size]
-    return port, additional_ports
+    # 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."
+        )
 
-def get_exception_traceback():
-    etype, value, tb = sys.exc_info()
-    err_str = "".join(traceback.format_exception(etype, value, tb))
-    return err_str
+    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)
@@ -178,14 +189,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
@@ -245,20 +253,104 @@ def wrap_kernel_launcher(kernel):
 
 
 def is_multimodal_model(model):
-    if isinstance(model, str):
-        return "llava" in model or "yi-vl" in model
     from sglang.srt.model_config import ModelConfig
 
+    if isinstance(model, str):
+        model = model.lower()
+        return "llava" in model or "yi-vl" in model or "llava-next" in model
+
     if isinstance(model, ModelConfig):
         model_path = model.path.lower()
-        return "llava" in model_path or "yi-vl" in model_path
-    raise Exception("unrecognized type")
+        return (
+            "llava" in model_path or "yi-vl" in model_path or "llava-next" in model_path
+        )
+
+    raise ValueError("unrecognized type")
+
+
+def decode_video_base64(video_base64):
+    from PIL import Image
+
+    # Decode the base64 string
+    video_bytes = base64.b64decode(video_base64)
+
+    # Placeholder for the start indices of each PNG image
+    img_starts = []
+
+    frame_format = "PNG"  # str(os.getenv('FRAME_FORMAT', "JPEG"))
+
+    assert frame_format in [
+        "PNG",
+        "JPEG",
+    ], "FRAME_FORMAT must be either 'PNG' or 'JPEG'"
+
+    if frame_format == "PNG":
+        # Find each PNG start signature to isolate images
+        i = 0
+        while i < len(video_bytes) - 7:  # Adjusted for the length of the PNG signature
+            # Check if we found the start of a PNG file
+            if (
+                video_bytes[i] == 0x89
+                and video_bytes[i + 1] == 0x50
+                and video_bytes[i + 2] == 0x4E
+                and video_bytes[i + 3] == 0x47
+                and video_bytes[i + 4] == 0x0D
+                and video_bytes[i + 5] == 0x0A
+                and video_bytes[i + 6] == 0x1A
+                and video_bytes[i + 7] == 0x0A
+            ):
+                img_starts.append(i)
+                i += 8  # Skip the PNG signature
+            else:
+                i += 1
+    else:
+        # Find each JPEG start (0xFFD8) to isolate images
+        i = 0
+        while (
+            i < len(video_bytes) - 1
+        ):  # Adjusted for the length of the JPEG SOI signature
+            # Check if we found the start of a JPEG file
+            if video_bytes[i] == 0xFF and video_bytes[i + 1] == 0xD8:
+                img_starts.append(i)
+                # Move to the next byte to continue searching for the next image start
+                i += 2
+            else:
+                i += 1
+
+    frames = []
+    for start_idx in img_starts:
+        # Assuming each image is back-to-back, the end of one image is the start of another
+        # The last image goes until the end of the byte string
+        end_idx = (
+            img_starts[img_starts.index(start_idx) + 1]
+            if img_starts.index(start_idx) + 1 < len(img_starts)
+            else len(video_bytes)
+        )
+        img_bytes = video_bytes[start_idx:end_idx]
+
+        # Convert bytes to a PIL Image
+        img = Image.open(BytesIO(img_bytes))
+
+        # Convert PIL Image to a NumPy array
+        frame = np.array(img)
+
+        # Append the frame to the list of frames
+        frames.append(frame)
+
+    # Ensure there's at least one frame to avoid errors with np.stack
+    if frames:
+        return np.stack(frames, axis=0), img.size
+    else:
+        return np.array([]), (
+            0,
+            0,
+        )  # Return an empty array and size tuple if no frames were found
 
 
 def load_image(image_file):
     from PIL import Image
 
-    image = None
+    image = image_size = None
 
     if image_file.startswith("http://") or image_file.startswith("https://"):
         timeout = int(os.getenv("REQUEST_TIMEOUT", "3"))
@@ -269,10 +361,71 @@ def load_image(image_file):
     elif image_file.startswith("data:"):
         image_file = image_file.split(",")[1]
         image = Image.open(BytesIO(base64.b64decode(image_file)))
+    elif image_file.startswith("video:"):
+        image_file = image_file.replace("video:", "")
+        image, image_size = decode_video_base64(image_file)
     else:
         image = Image.open(BytesIO(base64.b64decode(image_file)))
 
-    return image
+    return image, image_size
+
+
+def init_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 connect_to_rpyc_service(port, host="localhost"):
+    time.sleep(1)
+
+    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:
+            time.sleep(1)
+        repeat_count += 1
+    if repeat_count == 20:
+        raise RuntimeError("init rpc env error!")
+
+    return con.root
+
+
+def start_rpyc_process(service: rpyc.Service, port: int):
+    # Return the proxy and the process
+    proc = multiprocessing.Process(target=init_rpyc_service, args=(service, port))
+    proc.start()
+    proxy = connect_to_rpyc_service(port)
+    assert proc.is_alive()
+    return proxy, 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):
@@ -284,7 +437,30 @@ def assert_pkg_version(pkg: str, min_version: str):
                 f"is less than the minimum required version {min_version}"
             )
     except PackageNotFoundError:
-        raise Exception(f"{pkg} with minimum required version {min_version} is not installed")
+        raise Exception(
+            f"{pkg} with minimum required version {min_version} is not installed"
+        )
+
+
+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():
+    """
+    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.
+    """
+    import vllm.distributed.device_communicators.custom_all_reduce_utils as tgt
+
+    setattr(tgt, "gpu_p2p_access_check", lambda *arg, **kwargs: True)
 
 
 API_KEY_HEADER_NAME = "X-API-Key"
@@ -306,12 +482,3 @@ class APIKeyValidatorMiddleware(BaseHTTPMiddleware):
         response = await call_next(request)
         return response
 
-
-# FIXME: Remove this once we drop support for pydantic 1.x
-IS_PYDANTIC_1 = int(pydantic.VERSION.split(".")[0]) == 1
-
-
-def jsonify_pydantic_model(obj: BaseModel):
-    if IS_PYDANTIC_1:
-        return obj.json(ensure_ascii=False)
-    return obj.model_dump_json()

sglang/test/test_programs.py

@@ -304,6 +304,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 +350,46 @@ 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,40 +2,27 @@
 
 import base64
 import json
+import logging
+import signal
+import sys
 import threading
+import traceback
 import urllib.request
+from concurrent.futures import ThreadPoolExecutor
 from io import BytesIO
 from json import dumps
 
+import numpy as np
 import requests
 
 
-def get_available_gpu_memory(gpu_id, distributed=True):
-    """
-    Get available memory for cuda:gpu_id device.
-    When distributed is True, the available memory is the minimum available memory of all GPUs.
-    """
-    import torch
+logger = logging.getLogger(__name__)
 
-    num_gpus = torch.cuda.device_count()
-    assert gpu_id < num_gpus
 
-    if torch.cuda.current_device() != gpu_id:
-        print(
-            f"WARNING: current device is not {gpu_id}, but {torch.cuda.current_device()}, ",
-            "which may cause useless memory allocation for torch CUDA context.",
-        )
-
-    free_gpu_memory, _ = torch.cuda.mem_get_info(gpu_id)
-
-    if distributed:
-        tensor = torch.tensor(free_gpu_memory, dtype=torch.float32).to(
-            torch.device("cuda", gpu_id)
-        )
-        torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.MIN)
-        free_gpu_memory = tensor.item()
-
-    return free_gpu_memory / (1 << 30)
+def get_exception_traceback():
+    etype, value, tb = sys.exc_info()
+    err_str = "".join(traceback.format_exception(etype, value, tb))
+    return err_str
 
 
 def is_same_type(values):
@@ -110,8 +97,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):
@@ -130,6 +121,75 @@ def encode_image_base64(image_path):
         return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 
+def encode_frame(frame):
+    import cv2  # pip install opencv-python-headless
+    from PIL import Image
+
+    # Convert the frame to RGB (OpenCV uses BGR by default)
+    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+    # Convert the frame to PIL Image to easily convert to bytes
+    im_pil = Image.fromarray(frame)
+
+    # Convert to bytes
+    buffered = BytesIO()
+
+    # frame_format = str(os.getenv('FRAME_FORMAT', "JPEG"))
+
+    im_pil.save(buffered, format="PNG")
+
+    frame_bytes = buffered.getvalue()
+
+    # Return the bytes of the frame
+    return frame_bytes
+
+
+def encode_video_base64(video_path, num_frames=16):
+    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}")
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    print(f"target_frames: {num_frames}")
+
+    frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
+
+    frames = []
+    for i in range(total_frames):
+        ret, frame = cap.read()
+        if ret:
+            frames.append(frame)
+        else:
+            # Handle the case where the frame could not be read
+            # print(f"Warning: Could not read frame at index {i}.")
+            pass
+
+    cap.release()
+
+    # Safely select frames based on frame_indices, avoiding IndexError
+    frames = [frames[i] for i in frame_indices if i < len(frames)]
+
+    # If there are not enough frames, duplicate the last frame until we reach the target
+    while len(frames) < num_frames:
+        frames.append(frames[-1])
+
+    # Use ThreadPoolExecutor to process and encode frames in parallel
+    with ThreadPoolExecutor() as executor:
+        encoded_frames = list(executor.map(encode_frame, frames))
+
+    # encoded_frames = list(map(encode_frame, frames))
+
+    # Concatenate all frames bytes
+    video_bytes = b"".join(encoded_frames)
+
+    # Encode the concatenated bytes to base64
+    video_base64 = "video:" + base64.b64encode(video_bytes).decode("utf-8")
+
+    return video_base64
+
+
 def _is_chinese_char(cp):
     """Checks whether CP is the codepoint of a CJK character."""
     # This defines a "chinese character" as anything in the CJK Unicode block:
@@ -191,3 +251,12 @@ 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)