sglang 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

sglang/test/few_shot_gsm8k.py

@@ -0,0 +1,132 @@
+"""
+Run few-shot GSM-8K evaluation.
+
+Usage:
+python3 -m sglang.test.few_shot_gsm8k --num-questions 200
+"""
+
+import argparse
+import ast
+import re
+import time
+
+import numpy as np
+
+from sglang.api import set_default_backend
+from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
+from sglang.utils import download_and_cache_file, dump_state_text, read_jsonl
+
+INVALID = -9999999
+
+
+def get_one_example(lines, i, include_answer):
+    ret = "Question: " + lines[i]["question"] + "\nAnswer:"
+    if include_answer:
+        ret += " " + lines[i]["answer"]
+    return ret
+
+
+def get_few_shot_examples(lines, k):
+    ret = ""
+    for i in range(k):
+        ret += get_one_example(lines, i, True) + "\n\n"
+    return ret
+
+
+def get_answer_value(answer_str):
+    answer_str = answer_str.replace(",", "")
+    numbers = re.findall(r"\d+", answer_str)
+    if len(numbers) < 1:
+        return INVALID
+    try:
+        return ast.literal_eval(numbers[-1])
+    except SyntaxError:
+        return INVALID
+
+
+def main(args):
+    # Select backend
+    set_default_backend(RuntimeEndpoint(f"{args.host}:{args.port}"))
+
+    # Read data
+    url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl"
+    filename = download_and_cache_file(url)
+    lines = list(read_jsonl(filename))
+
+    # Construct prompts
+    num_questions = args.num_questions
+    num_shots = args.num_shots
+    few_shot_examples = get_few_shot_examples(lines, num_shots)
+
+    questions = []
+    labels = []
+    for i in range(len(lines[:num_questions])):
+        questions.append(get_one_example(lines, i, False))
+        labels.append(get_answer_value(lines[i]["answer"]))
+    assert all(l != INVALID for l in labels)
+    arguments = [{"question": q} for q in questions]
+
+    #####################################
+    ######### SGL Program Begin #########
+    #####################################
+
+    import sglang as sgl
+
+    @sgl.function
+    def few_shot_gsm8k(s, question):
+        s += few_shot_examples + question
+        s += sgl.gen(
+            "answer", max_tokens=512, stop=["Question", "Assistant:", "<|separator|>"]
+        )
+
+    #####################################
+    ########## SGL Program End ##########
+    #####################################
+
+    # Run requests
+    tic = time.time()
+    states = few_shot_gsm8k.run_batch(
+        arguments,
+        temperature=0,
+        num_threads=args.parallel,
+        progress_bar=True,
+    )
+    latency = time.time() - tic
+
+    preds = []
+    for i in range(len(states)):
+        preds.append(get_answer_value(states[i]["answer"]))
+
+    # print(f"{preds=}")
+    # print(f"{labels=}")
+
+    # Compute accuracy
+    acc = np.mean(np.array(preds) == np.array(labels))
+    invalid = np.mean(np.array(preds) == INVALID)
+
+    # Compute speed
+    num_output_tokens = sum(
+        s.get_meta_info("answer")["completion_tokens"] for s in states
+    )
+    output_throughput = num_output_tokens / latency
+
+    # Print results
+    print(f"Accuracy: {acc:.3f}")
+    print(f"Invalid: {invalid:.3f}")
+    print(f"Latency: {latency:.3f} s")
+    print(f"Output throughput: {output_throughput:.3f} token/s")
+
+    # Dump results
+    dump_state_text("tmp_output_gsm8k.txt", states)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-shots", type=int, default=5)
+    parser.add_argument("--data-path", type=str, default="test.jsonl")
+    parser.add_argument("--num-questions", type=int, default=200)
+    parser.add_argument("--parallel", type=int, default=128)
+    parser.add_argument("--host", type=str, default="http://127.0.0.1")
+    parser.add_argument("--port", type=int, default=30000)
+    args = parser.parse_args()
+    main(args)

sglang/test/runners.py

@@ -21,6 +21,7 @@ from typing import List, Union
 
 import torch
 import torch.nn.functional as F
+from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from sglang.srt.server import Runtime
@@ -50,6 +51,13 @@ def get_dtype_str(torch_dtype):
         raise NotImplementedError()
 
 
+def get_top_logprobs(logits, k):
+    logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
+    del logits
+    logprobs, top_indices = torch.topk(logprobs, k=k, dim=-1)
+    return logprobs
+
+
 @dataclass
 class ModelOutput:
     output_strs: List[str] = None
@@ -65,8 +73,10 @@ class HFRunner:
         model_path,
         torch_dtype,
         is_generation,
+        output_str_only=False,
     ):
         self.is_generation = is_generation
+        self.output_str_only = output_str_only
 
         self.in_queue = mp.Queue()
         self.out_queue = mp.Queue()
@@ -89,7 +99,7 @@ class HFRunner:
         )
 
         if self.is_generation:
-            self.model = AutoModelForCausalLM.from_pretrained(
+            self.base_model = AutoModelForCausalLM.from_pretrained(
                 model_path,
                 torch_dtype=torch_dtype,
                 trust_remote_code=False,
@@ -104,12 +114,16 @@ class HFRunner:
             )
 
         while True:
-            prompts, max_new_tokens = in_queue.get()
+            prompts, max_new_tokens, lora_paths = in_queue.get()
+            if lora_paths is not None:
+                assert len(prompts) == len(lora_paths)
+
             if prompts is not None:
                 if self.is_generation:
                     output_strs = []
-                    prefill_logprobs = []
-                    for p in prompts:
+                    top_input_logprobs = []
+                    top_output_logprobs = []
+                    for i, p in enumerate(prompts):
                         if isinstance(p, str):
                             input_ids = self.tokenizer.encode(
                                 p, return_tensors="pt"
@@ -117,40 +131,68 @@ class HFRunner:
                         else:
                             input_ids = torch.tensor([p], device="cuda")
 
-                        output_ids = self.model.generate(
-                            input_ids, do_sample=False, max_new_tokens=max_new_tokens
+                        if lora_paths is not None and lora_paths[i] is not None:
+                            self.model = PeftModel.from_pretrained(
+                                self.base_model,
+                                lora_paths[i],
+                                torch_dtype=torch_dtype,
+                                is_trainable=False,
+                            )
+                        else:
+                            self.model = self.base_model
+
+                        outputs = self.model.generate(
+                            input_ids,
+                            do_sample=False,
+                            temperature=None,
+                            top_p=None,
+                            max_new_tokens=max_new_tokens,
+                            return_dict_in_generate=True,
+                            output_scores=(not self.output_str_only),
                         )
                         output_strs.append(
-                            self.tokenizer.decode(output_ids[0][len(input_ids[0]) :])
+                            self.tokenizer.decode(outputs[0][0][len(input_ids[0]) :])
                         )
-
-                        logits = self.model.forward(input_ids).logits[0]
-                        logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
-                        logprobs, top_indices = torch.topk(
-                            logprobs, k=NUM_TOP_LOGPROBS, dim=-1
-                        )
-                        # print("index", top_indices)
-                        prefill_logprobs.append(logprobs.tolist())
-                        del logits
-                        del logprobs
+                        if not self.output_str_only:
+                            # outputs.scores: (num_token, 1, vocab_size)
+                            top_output_logprobs.append(
+                                [
+                                    get_top_logprobs(
+                                        logits[0], NUM_TOP_LOGPROBS
+                                    ).tolist()
+                                    for logits in outputs.scores
+                                ]
+                            )
+                            del outputs
+
+                            input_logits = self.model.forward(input_ids).logits[0]
+                            top_input_logprobs.append(
+                                get_top_logprobs(
+                                    input_logits, NUM_TOP_LOGPROBS
+                                ).tolist()
+                            )
+                            del input_logits
 
                     out_queue.put(
                         ModelOutput(
-                            output_strs=output_strs, top_input_logprobs=prefill_logprobs
+                            output_strs=output_strs,
+                            top_input_logprobs=top_input_logprobs,
+                            top_output_logprobs=top_output_logprobs,
                         )
                     )
 
                 else:
+                    assert not self.output_str_only
                     logits = self.model.encode(prompts).tolist()
-
                     out_queue.put(ModelOutput(embed_logits=logits))
 
     def forward(
         self,
         prompts: Union[List[str], List[torch.Tensor]] = DEFAULT_PROMPTS,
         max_new_tokens=8,
+        lora_paths=None,
     ):
-        self.in_queue.put((prompts, max_new_tokens))
+        self.in_queue.put((prompts, max_new_tokens, lora_paths))
         return self.out_queue.get()
 
     def terminate(self):
@@ -173,6 +215,10 @@ class SRTRunner:
         is_generation,
         tp_size=1,
         port=DEFAULT_PORT_FOR_SRT_TEST_RUNNER,
+        lora_paths=None,
+        max_loras_per_batch=4,
+        disable_cuda_graph=False,
+        disable_radix_cache=False,
     ):
         self.is_generation = is_generation
         self.runtime = Runtime(
@@ -183,21 +229,28 @@ class SRTRunner:
             mem_fraction_static=0.69,
             trust_remote_code=False,
             is_embedding=not self.is_generation,
+            lora_paths=lora_paths,
+            max_loras_per_batch=max_loras_per_batch,
+            disable_cuda_graph=disable_cuda_graph,
+            disable_radix_cache=disable_radix_cache,
         )
 
     def forward(
         self,
         prompts: Union[List[str], List[torch.Tensor]] = DEFAULT_PROMPTS,
         max_new_tokens=8,
+        lora_paths=None,
     ):
         if self.is_generation:
             # the return value contains logprobs from prefill
             output_strs = []
             top_input_logprobs = []
+            top_output_logprobs = []
             sampling_params = {"max_new_tokens": max_new_tokens, "temperature": 0}
-            for prompt in prompts:
+            for i, prompt in enumerate(prompts):
                 response = self.runtime.generate(
                     prompt,
+                    lora_path=lora_paths[i] if lora_paths else None,
                     sampling_params=sampling_params,
                     return_logprob=True,
                     logprob_start_len=0,
@@ -219,9 +272,48 @@ class SRTRunner:
                         ]
                     ]
                 )
+                top_output_logprobs.append(
+                    [
+                        [tup[0] for tup in x[:NUM_TOP_LOGPROBS]]
+                        for x in response["meta_info"]["output_top_logprobs"]
+                    ]
+                )
+
+            return ModelOutput(
+                output_strs=output_strs,
+                top_input_logprobs=top_input_logprobs,
+                top_output_logprobs=top_output_logprobs,
+            )
+        else:
+            response = self.runtime.encode(prompts)
+            response = json.loads(response)
+            logits = [x["embedding"] for x in response]
+            return ModelOutput(embed_logits=logits)
+
+    def batch_forward(
+        self,
+        prompts: Union[List[str], List[torch.Tensor]] = DEFAULT_PROMPTS,
+        max_new_tokens=8,
+        lora_paths=None,
+    ):
+        """
+        testing serving by sending all prompts once
+        only return output strings and no logprobs
+        """
+        if self.is_generation:
+            # the return value contains logprobs from prefill
+            output_strs = []
+            sampling_params = {"max_new_tokens": max_new_tokens, "temperature": 0}
+            response = self.runtime.generate(
+                prompts,
+                lora_path=lora_paths if lora_paths else None,
+                sampling_params=sampling_params,
+            )
+            response = json.loads(response)
+            output_strs = [r["text"] for r in response]
 
             return ModelOutput(
-                output_strs=output_strs, top_input_logprobs=top_input_logprobs
+                output_strs=output_strs,
             )
         else:
             response = self.runtime.encode(prompts)

sglang/test/test_programs.py

@@ -7,7 +7,7 @@ import time
 import numpy as np
 
 import sglang as sgl
-from sglang.utils import fetch_and_cache_jsonl
+from sglang.utils import download_and_cache_file, read_jsonl
 
 
 def test_few_shot_qa():
@@ -456,10 +456,6 @@ def test_chat_completion_speculative():
 def test_hellaswag_select():
     """Benchmark the accuracy of sgl.select on the HellaSwag dataset."""
 
-    url = "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl"
-    lines = fetch_and_cache_jsonl(url)
-
-    # Construct prompts
     def get_one_example(lines, i, include_answer):
         ret = lines[i]["activity_label"] + ": " + lines[i]["ctx"] + " "
         if include_answer:
@@ -472,6 +468,12 @@ def test_hellaswag_select():
             ret += get_one_example(lines, i, True) + "\n\n"
         return ret
 
+    # Read data
+    url = "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl"
+    filename = download_and_cache_file(url)
+    lines = list(read_jsonl(filename))
+
+    # Construct prompts
     num_questions = 200
     num_shots = 20
     few_shot_examples = get_few_shot_examples(lines, num_shots)

sglang/test/test_utils.py

@@ -7,6 +7,7 @@ import subprocess
 import threading
 import time
 from functools import partial
+from types import SimpleNamespace
 from typing import Callable, List, Optional
 
 import numpy as np
@@ -14,6 +15,7 @@ import requests
 import torch
 import torch.nn.functional as F
 
+from sglang.bench_serving import run_benchmark
 from sglang.global_config import global_config
 from sglang.lang.backend.openai import OpenAI
 from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
@@ -28,7 +30,13 @@ DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2 = "meta-llama/Meta-Llama-3.1-70B-Instruc
 DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP1 = "neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8,neuralmagic/Mistral-7B-Instruct-v0.3-FP8,neuralmagic/DeepSeek-Coder-V2-Lite-Instruct-FP8,neuralmagic/gemma-2-2b-it-FP8"
 DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2 = "neuralmagic/Meta-Llama-3.1-70B-Instruct-FP8,neuralmagic/Mixtral-8x7B-Instruct-v0.1-FP8,neuralmagic/Qwen2-72B-Instruct-FP8,neuralmagic/Qwen2-57B-A14B-Instruct-FP8"
 
-if os.getenv("SGLANG_IS_IN_CI", "false") == "true":
+
+def is_in_ci():
+    """Return whether it is in CI runner."""
+    return os.getenv("SGLANG_IS_IN_CI", "false") == "true"
+
+
+if is_in_ci():
     DEFAULT_PORT_FOR_SRT_TEST_RUNNER = 5157
     DEFAULT_URL_FOR_TEST = "http://127.0.0.1:6157"
 else:
@@ -501,3 +509,79 @@ def run_unittest_files(files: List[str], timeout_per_file: float):
 
 def get_similarities(vec1, vec2):
     return F.cosine_similarity(torch.tensor(vec1), torch.tensor(vec2), dim=0)
+
+
+def run_bench_serving(model, num_prompts, request_rate, other_server_args):
+    # Launch the server
+    base_url = DEFAULT_URL_FOR_TEST
+    process = popen_launch_server(
+        model,
+        base_url,
+        timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+        other_args=other_server_args,
+    )
+
+    # Run benchmark
+    args = SimpleNamespace(
+        backend="sglang",
+        base_url=base_url,
+        host=None,
+        port=None,
+        dataset_name="random",
+        dataset_path="",
+        model=None,
+        tokenizer=None,
+        num_prompts=num_prompts,
+        sharegpt_output_len=None,
+        random_input_len=4096,
+        random_output_len=2048,
+        random_range_ratio=0.0,
+        request_rate=request_rate,
+        multi=None,
+        seed=0,
+        output_file=None,
+        disable_tqdm=False,
+        disable_stream=False,
+        disable_ignore_eos=False,
+        extra_request_body=None,
+    )
+
+    try:
+        res = run_benchmark(args)
+    finally:
+        kill_child_process(process.pid)
+
+    assert res["completed"] == num_prompts
+    return res
+
+
+def run_bench_latency(model, other_args):
+    command = [
+        "python3",
+        "-m",
+        "sglang.bench_latency",
+        "--model-path",
+        model,
+        "--batch-size",
+        "1",
+        "--input",
+        "128",
+        "--output",
+        "8",
+        *other_args,
+    ]
+    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+    try:
+        stdout, stderr = process.communicate()
+        output = stdout.decode()
+        error = stderr.decode()
+        print(f"Output: {output}", flush=True)
+        print(f"Error: {error}", flush=True)
+
+        lastline = output.split("\n")[-3]
+        output_throughput = float(lastline.split(" ")[-2])
+    finally:
+        kill_child_process(process.pid)
+
+    return output_throughput

sglang/utils.py

@@ -12,7 +12,7 @@ import urllib.request
 from concurrent.futures import ThreadPoolExecutor
 from io import BytesIO
 from json import dumps
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 import requests
@@ -38,13 +38,11 @@ def is_same_type(values: list):
 
 def read_jsonl(filename: str):
     """Read a JSONL file."""
-    rets = []
     with open(filename) as fin:
         for line in fin:
             if line.startswith("#"):
                 continue
-            rets.append(json.loads(line))
-    return rets
+            yield json.loads(line)
 
 
 def dump_state_text(filename: str, states: list, mode: str = "w"):
@@ -264,38 +262,35 @@ class LazyImport:
         return module(*args, **kwargs)
 
 
-def fetch_and_cache_jsonl(url, cache_file="cached_data.jsonl"):
-    """Read and cache a jsonl file from a url."""
+def download_and_cache_file(url: str, filename: Optional[str] = None):
+    """Read and cache a file from a url."""
+    if filename is None:
+        filename = os.path.join("/tmp", url.split("/")[-1])
 
     # Check if the cache file already exists
-    if os.path.exists(cache_file):
-        print("Loading data from cache...")
-        with open(cache_file, "r") as f:
-            data = [json.loads(line) for line in f]
-    else:
-        print("Downloading data from URL...")
-        # Stream the response to show the progress bar
-        response = requests.get(url, stream=True)
-        response.raise_for_status()  # Check for request errors
-
-        # Total size of the file in bytes
-        total_size = int(response.headers.get("content-length", 0))
-        chunk_size = 1024  # Download in chunks of 1KB
-
-        # Use tqdm to display the progress bar
-        with open(cache_file, "wb") as f, tqdm(
-            desc=cache_file,
-            total=total_size,
-            unit="B",
-            unit_scale=True,
-            unit_divisor=1024,
-        ) as bar:
-            for chunk in response.iter_content(chunk_size=chunk_size):
-                f.write(chunk)
-                bar.update(len(chunk))
-
-        # Convert the data to a list of dictionaries
-        with open(cache_file, "r") as f:
-            data = [json.loads(line) for line in f]
-
-    return data
+    if os.path.exists(filename):
+        return filename
+
+    print(f"Downloading from {url} to {filename}")
+
+    # Stream the response to show the progress bar
+    response = requests.get(url, stream=True)
+    response.raise_for_status()  # Check for request errors
+
+    # Total size of the file in bytes
+    total_size = int(response.headers.get("content-length", 0))
+    chunk_size = 1024  # Download in chunks of 1KB
+
+    # Use tqdm to display the progress bar
+    with open(filename, "wb") as f, tqdm(
+        desc=filename,
+        total=total_size,
+        unit="B",
+        unit_scale=True,
+        unit_divisor=1024,
+    ) as bar:
+        for chunk in response.iter_content(chunk_size=chunk_size):
+            f.write(chunk)
+            bar.update(len(chunk))
+
+    return filename

sglang/version.py

@@ -1 +1 @@
-__version__ = "0.3.0"
+__version__ = "0.3.1"