sglang 0.1.14__py3-none-any.whl → 0.1.21__py3-none-any.whl

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
@@ -226,7 +224,7 @@ Action 3: Finish [United States].\n
 
 def test_parallel_decoding():
     max_tokens = 64
-    number = 5
+    fork_size = 5
 
     @sgl.function
     def parallel_decoding(s, topic):
@@ -234,17 +232,17 @@ def test_parallel_decoding():
         s += "USER: Give some tips for " + topic + ".\n"
         s += (
             "ASSISTANT: Okay. Here are "
-            + str(number)
+            + str(fork_size)
             + " concise tips, each under 8 words:\n"
         )
 
         # Generate skeleton
-        for i in range(1, 1 + number):
+        for i in range(1, 1 + fork_size):
             s += f"{i}." + sgl.gen(max_tokens=16, stop=[".", "\n"]) + ".\n"
 
         # Generate detailed tips
-        forks = s.fork(number)
-        for i in range(number):
+        forks = s.fork(fork_size)
+        for i in range(fork_size):
             forks[
                 i
             ] += f"Now, I expand tip {i+1} into a detailed paragraph:\nTip {i+1}:"
@@ -253,7 +251,7 @@ def test_parallel_decoding():
 
         # Concatenate tips and summarize
         s += "Here are these tips with detailed explanation:\n"
-        for i in range(number):
+        for i in range(fork_size):
             s += f"Tip {i+1}:" + forks[i]["detailed_tip"] + "\n"
 
         s += "\nIn summary," + sgl.gen("summary", max_tokens=512)
@@ -296,7 +294,7 @@ def test_parallel_encoding(check_answer=True):
 def test_image_qa():
     @sgl.function
     def image_qa(s, question):
-        s += sgl.user(sgl.image("test_image.png") + question)
+        s += sgl.user(sgl.image("example_image.png") + question)
         s += sgl.assistant(sgl.gen("answer"))
 
     state = image_qa.run(
@@ -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"]
@@ -313,6 +312,7 @@ def test_image_qa():
 def test_stream():
     @sgl.function
     def qa(s, question):
+        s += sgl.system("You are a helpful assistant.")
         s += sgl.user(question)
         s += sgl.assistant(sgl.gen("answer"))
 
@@ -348,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

@@ -1,13 +1,20 @@
 """Common utilities for testing and benchmarking"""
 
+import asyncio
+from functools import partial
+
 import numpy as np
 import requests
+
 from sglang.backend.openai import OpenAI
 from sglang.backend.runtime_endpoint import RuntimeEndpoint
 from sglang.global_config import global_config
+from sglang.utils import get_exception_traceback
 
 
-def call_generate_lightllm(prompt, temperature, max_tokens, stop, url):
+def call_generate_lightllm(prompt, temperature, max_tokens, stop=None, url=None):
+    assert url is not None
+
     data = {
         "inputs": prompt,
         "parameters": {
@@ -22,7 +29,9 @@ def call_generate_lightllm(prompt, temperature, max_tokens, stop, url):
     return pred
 
 
-def call_generate_vllm(prompt, temperature, max_tokens, stop, url, n=1):
+def call_generate_vllm(prompt, temperature, max_tokens, stop=None, n=1, url=None):
+    assert url is not None
+
     data = {
         "prompt": prompt,
         "temperature": temperature,
@@ -40,8 +49,10 @@ def call_generate_vllm(prompt, temperature, max_tokens, stop, url, n=1):
 
 
 def call_generate_outlines(
-    prompt, temperature, max_tokens, url, stop=[], regex=None, n=1
+    prompt, temperature, max_tokens, stop=[], regex=None, n=1, url=None
 ):
+    assert url is not None
+
     data = {
         "prompt": prompt,
         "temperature": temperature,
@@ -59,7 +70,9 @@ def call_generate_outlines(
     return pred
 
 
-def call_generate_srt_raw(prompt, temperature, max_tokens, stop, url):
+def call_generate_srt_raw(prompt, temperature, max_tokens, stop=None, url=None):
+    assert url is not None
+
     data = {
         "text": prompt,
         "sampling_params": {
@@ -75,7 +88,98 @@ def call_generate_srt_raw(prompt, temperature, max_tokens, stop, url):
     return pred
 
 
-def call_select_lightllm(context, choices, url):
+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
+):
+    assert model is not None
+    from guidance import gen
+
+    rets = []
+    for _ in range(n):
+        out = (
+            model
+            + prompt
+            + gen(
+                name="answer",
+                max_tokens=max_tokens,
+                temperature=temperature,
+                stop=stop,
+                regex=regex,
+            )
+        )
+        rets.append(out["answer"])
+    return rets if n > 1 else rets[0]
+
+
+async def call_generate_lmql(
+    prompt, temperature, max_tokens, stop=None, n=1, max_len=4096, model=None, **kwargs
+):
+    assert model is not None
+    import lmql
+
+    if stop != None:
+
+        @lmql.query(model=model)
+        async def program(question, max_tokens, stop):
+            '''lmql
+            """{question}[ANSWER]""" where len(TOKENS(ANSWER)) < max_tokens and STOPS_AT(ANSWER, stop)
+            return ANSWER
+            '''
+
+    else:
+
+        @lmql.query(model=model)
+        async def program(question, max_tokens):
+            '''lmql
+            """{question}[ANSWER]""" where len(TOKENS(ANSWER)) < max_tokens
+            return ANSWER
+            '''
+
+    tasks = [
+        program(
+            question=prompt,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            stop=stop,
+            max_len=max_len,
+            **kwargs,
+        )
+        for _ in range(n)
+    ]
+    rets = await asyncio.gather(*tasks)
+    return rets if n > 1 else rets[0]
+
+
+def call_select_lightllm(context, choices, url=None):
+    assert url is not None
+
     scores = []
     for i in range(len(choices)):
         data = {
@@ -90,7 +194,9 @@ def call_select_lightllm(context, choices, url):
     return np.argmax(scores)
 
 
-def call_select_vllm(context, choices, url):
+def call_select_vllm(context, choices, url=None):
+    assert url is not None
+
     scores = []
     for i in range(len(choices)):
         data = {
@@ -112,6 +218,31 @@ def call_select_vllm(context, choices, url):
     """
 
 
+def call_select_guidance(context, choices, model=None):
+    assert model is not None
+    from guidance import select
+
+    out = model + context + select(choices, name="answer")
+    return choices.index(out["answer"])
+
+
+async def call_select_lmql(context, choices, temperature=0, max_len=4096, model=None):
+    assert model is not None
+    import lmql
+
+    @lmql.query(model=model)
+    async def program(ctx, choices):
+        '''lmql
+        """{ctx}[ANSWER]""" where ANSWER in set(choices)
+        return ANSWER
+        '''
+
+    answer = await program(
+        ctx=context, choices=choices, temperature=temperature, max_len=max_len
+    )
+    return choices.index(answer)
+
+
 def add_common_other_args_and_parse(parser):
     parser.add_argument("--parallel", type=int, default=64)
     parser.add_argument("--host", type=str, default="http://127.0.0.1")
@@ -120,8 +251,18 @@ def add_common_other_args_and_parse(parser):
         "--backend",
         type=str,
         required=True,
-        choices=["vllm", "lightllm", "guidance", "lmql", "srt-raw", "llama.cpp"],
+        choices=[
+            "vllm",
+            "outlines",
+            "lightllm",
+            "ginfer",
+            "guidance",
+            "lmql",
+            "srt-raw",
+            "llama.cpp",
+        ],
     )
+    parser.add_argument("--n-ctx", type=int, default=4096)
     parser.add_argument(
         "--model-path", type=str, default="meta-llama/Llama-2-7b-chat-hf"
     )
@@ -131,9 +272,11 @@ def add_common_other_args_and_parse(parser):
     if args.port is None:
         default_port = {
             "vllm": 21000,
+            "outlines": 21000,
             "lightllm": 22000,
             "lmql": 23000,
             "srt-raw": 30000,
+            "ginfer": 9988,
         }
         args.port = default_port.get(args.backend, None)
     return args
@@ -160,3 +303,79 @@ def select_sglang_backend(args):
     else:
         raise ValueError(f"Invalid backend: {args.backend}")
     return backend
+
+
+def _get_call_generate(args):
+    if args.backend == "lightllm":
+        return partial(call_generate_lightllm, url=f"{args.host}:{args.port}/generate")
+    elif args.backend == "vllm":
+        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":
+        from guidance import models
+
+        model = models.LlamaCpp(args.model_path, n_gpu_layers=-1, n_ctx=args.n_ctx)
+        call_generate = partial(call_generate_guidance, model=model)
+        call_generate("Hello,", 1.0, 8, ".")
+        return call_generate
+    elif args.backend == "lmql":
+        import lmql
+
+        model = lmql.model(args.model_path, endpoint=f"{args.host}:{args.port}")
+        return partial(call_generate_lmql, model=model)
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+
+def _get_call_select(args):
+    if args.backend == "lightllm":
+        return partial(call_select_lightllm, url=f"{args.host}:{args.port}/generate")
+    elif args.backend == "vllm":
+        return partial(call_select_vllm, url=f"{args.host}:{args.port}/generate")
+    elif args.backend == "guidance":
+        from guidance import models
+
+        model = models.LlamaCpp(args.model_path, n_gpu_layers=-1, n_ctx=args.n_ctx)
+        call_select = partial(call_select_guidance, model=model)
+
+        call_select("Hello,", ["world", "earth"])
+        return call_select
+
+    elif args.backend == "lmql":
+        import lmql
+
+        model = lmql.model(args.model_path, endpoint=f"{args.host}:{args.port}")
+        return partial(call_select_lmql, model=model)
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+
+def get_call_generate(args):
+    call_generate = _get_call_generate(args)
+
+    def func(*args, **kwargs):
+        try:
+            return call_generate(*args, **kwargs)
+        except Exception:
+            print("Exception in call_generate:\n" + get_exception_traceback())
+            raise
+
+    return func
+
+
+def get_call_select(args):
+    call_select = _get_call_select(args)
+
+    def func(*args, **kwargs):
+        try:
+            return call_select(*args, **kwargs)
+        except Exception:
+            print("Exception in call_select:\n" + get_exception_traceback())
+            raise
+
+    return func

sglang/utils.py

@@ -2,40 +2,26 @@
 
 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
 
+logger = logging.getLogger(__name__)
 
-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
 
-    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 +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):
@@ -130,6 +120,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 +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)