test_utils.py 29.8 KB
Newer Older
Lianmin Zheng's avatar
Lianmin Zheng committed
1
"""Common utilities for testing and benchmarking"""
2

3
import argparse
4
import copy
5
import logging
6
import os
7
import random
8
import subprocess
9
import threading
10
import time
11
import traceback
12
import unittest
13
from concurrent.futures import ThreadPoolExecutor
Byron Hsu's avatar
Byron Hsu committed
14
from dataclasses import dataclass
Liangsheng Yin's avatar
Liangsheng Yin committed
15
from functools import partial
16
from types import SimpleNamespace
17
from typing import Callable, List, Optional, Tuple
Liangsheng Yin's avatar
Liangsheng Yin committed
18

Lianmin Zheng's avatar
Lianmin Zheng committed
19
20
import numpy as np
import requests
21
22
import torch
import torch.nn.functional as F
Liangsheng Yin's avatar
Liangsheng Yin committed
23

24
from sglang.bench_serving import run_benchmark
Lianmin Zheng's avatar
Lianmin Zheng committed
25
from sglang.global_config import global_config
Ying Sheng's avatar
Ying Sheng committed
26
27
from sglang.lang.backend.openai import OpenAI
from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
28
29
30
31
32
33
from sglang.srt.utils import (
    get_bool_env_var,
    is_port_available,
    kill_process_tree,
    retry,
)
34
from sglang.test.run_eval import run_eval
35
from sglang.utils import get_exception_traceback
Liangsheng Yin's avatar
Liangsheng Yin committed
36

Lianmin Zheng's avatar
Lianmin Zheng committed
37
DEFAULT_FP8_MODEL_NAME_FOR_TEST = "neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8"
HandH1998's avatar
HandH1998 committed
38
39
40
41
DEFAULT_FP8_MODEL_NAME_FOR_ACCURACY_TEST = "neuralmagic/Meta-Llama-3-8B-Instruct-FP8"
DEFAULT_FP8_MODEL_NAME_FOR_DYNAMIC_QUANT_ACCURACY_TEST = (
    "neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic"
)
42
43
44
45
DEFAULT_FP8_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST = (
    "nvidia/Llama-3.1-8B-Instruct-FP8"
)

46
DEFAULT_MODEL_NAME_FOR_TEST = "meta-llama/Llama-3.1-8B-Instruct"
Lianmin Zheng's avatar
Lianmin Zheng committed
47
DEFAULT_SMALL_MODEL_NAME_FOR_TEST = "meta-llama/Llama-3.2-1B-Instruct"
Yineng Zhang's avatar
Yineng Zhang committed
48
DEFAULT_MOE_MODEL_NAME_FOR_TEST = "mistralai/Mixtral-8x7B-Instruct-v0.1"
49
50
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST = "Qwen/Qwen1.5-MoE-A2.7B"
DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST = "Alibaba-NLP/gte-Qwen2-1.5B-instruct"
Ke Bao's avatar
Ke Bao committed
51
DEFAULT_MLA_MODEL_NAME_FOR_TEST = "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"
Yineng Zhang's avatar
Yineng Zhang committed
52
DEFAULT_MLA_FP8_MODEL_NAME_FOR_TEST = "neuralmagic/DeepSeek-Coder-V2-Lite-Instruct-FP8"
Xihuai Wang's avatar
Xihuai Wang committed
53
DEFAULT_REASONING_MODEL_NAME_FOR_TEST = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
54
55
56
DEFAULT_AWQ_MOE_MODEL_NAME_FOR_TEST = (
    "hugging-quants/Mixtral-8x7B-Instruct-v0.1-AWQ-INT4"
)
57
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH = 1000
58
DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1 = "meta-llama/Llama-3.1-8B-Instruct,mistralai/Mistral-7B-Instruct-v0.3,deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct,google/gemma-2-27b-it"
59
DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2 = "meta-llama/Llama-3.1-70B-Instruct,mistralai/Mixtral-8x7B-Instruct-v0.1,Qwen/Qwen2-57B-A14B-Instruct"
60
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"
Ke Bao's avatar
Ke Bao committed
61
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,neuralmagic/DeepSeek-Coder-V2-Lite-Instruct-FP8"
62
DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_QUANT_TP1 = "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4,hugging-quants/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4,hugging-quants/Mixtral-8x7B-Instruct-v0.1-AWQ-INT4"
63
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN = "Qwen/Qwen2.5-1.5B-Instruct"
64
65
DEFAULT_SMALL_VLM_MODEL_NAME = "Qwen/Qwen2-VL-2B"

66
DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST = "meta-llama/Llama-2-7b-chat-hf"
67
DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST = "lmsys/sglang-EAGLE-llama2-chat-7B"
68

69
70
71
DEFAULT_IMAGE_URL = "https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true"
DEFAULT_VIDEO_URL = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4"

72
73
74

def is_in_ci():
    """Return whether it is in CI runner."""
75
    return get_bool_env_var("SGLANG_IS_IN_CI")
76
77
78


if is_in_ci():
79
80
81
    DEFAULT_PORT_FOR_SRT_TEST_RUNNER = (
        5000 + int(os.environ.get("CUDA_VISIBLE_DEVICES", "0")[0]) * 100
    )
82
else:
83
84
85
86
    DEFAULT_PORT_FOR_SRT_TEST_RUNNER = (
        7000 + int(os.environ.get("CUDA_VISIBLE_DEVICES", "0")[0]) * 100
    )
DEFAULT_URL_FOR_TEST = f"http://127.0.0.1:{DEFAULT_PORT_FOR_SRT_TEST_RUNNER + 1000}"
87

Lianmin Zheng's avatar
Lianmin Zheng committed
88

Liangsheng Yin's avatar
Liangsheng Yin committed
89
90
def call_generate_lightllm(prompt, temperature, max_tokens, stop=None, url=None):
    assert url is not None
Lianmin Zheng's avatar
Lianmin Zheng committed
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105

    data = {
        "inputs": prompt,
        "parameters": {
            "temperature": temperature,
            "max_new_tokens": max_tokens,
            "stop_sequences": stop,
        },
    }
    res = requests.post(url, json=data)
    assert res.status_code == 200
    pred = res.json()["generated_text"][0]
    return pred


106
107
108
109
110
111
112
113
114
115
116
def find_available_port(base_port: int):
    port = base_port + random.randint(100, 1000)
    while True:
        if is_port_available(port):
            return port
        if port < 60000:
            port += 42
        else:
            port -= 43


Liangsheng Yin's avatar
Liangsheng Yin committed
117
118
119
def call_generate_vllm(prompt, temperature, max_tokens, stop=None, n=1, url=None):
    assert url is not None

Lianmin Zheng's avatar
Lianmin Zheng committed
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
    data = {
        "prompt": prompt,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stop": stop,
        "n": n,
    }
    res = requests.post(url, json=data)
    assert res.status_code == 200
    if n == 1:
        pred = res.json()["text"][0][len(prompt) :]
    else:
        pred = [x[len(prompt) :] for x in res.json()["text"]]
    return pred


136
def call_generate_outlines(
137
    prompt, temperature, max_tokens, stop=None, regex=None, n=1, url=None
138
):
Liangsheng Yin's avatar
Liangsheng Yin committed
139
140
    assert url is not None

141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
    data = {
        "prompt": prompt,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stop": stop,
        "regex": regex,
        "n": n,
    }
    res = requests.post(url, json=data)
    assert res.status_code == 200
    if n == 1:
        pred = res.json()["text"][0][len(prompt) :]
    else:
        pred = [x[len(prompt) :] for x in res.json()["text"]]
    return pred


Liangsheng Yin's avatar
Liangsheng Yin committed
158
159
160
def call_generate_srt_raw(prompt, temperature, max_tokens, stop=None, url=None):
    assert url is not None

Lianmin Zheng's avatar
Lianmin Zheng committed
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
    data = {
        "text": prompt,
        "sampling_params": {
            "temperature": temperature,
            "max_new_tokens": max_tokens,
            "stop": stop,
        },
    }
    res = requests.post(url, json=data)
    assert res.status_code == 200
    obj = res.json()
    pred = obj["text"]
    return pred


Liangsheng Yin's avatar
Liangsheng Yin committed
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
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]


def call_select_lightllm(context, choices, url=None):
    assert url is not None

Lianmin Zheng's avatar
Lianmin Zheng committed
202
203
204
205
206
207
208
209
210
211
212
213
214
215
    scores = []
    for i in range(len(choices)):
        data = {
            "inputs": context + choices[i],
            "parameters": {
                "max_new_tokens": 1,
            },
        }
        res = requests.post(url, json=data)
        assert res.status_code == 200
        scores.append(0)
    return np.argmax(scores)


Liangsheng Yin's avatar
Liangsheng Yin committed
216
217
218
def call_select_vllm(context, choices, url=None):
    assert url is not None

Lianmin Zheng's avatar
Lianmin Zheng committed
219
220
221
222
223
224
225
226
227
    scores = []
    for i in range(len(choices)):
        data = {
            "prompt": context + choices[i],
            "max_tokens": 1,
            "prompt_logprobs": 1,
        }
        res = requests.post(url, json=data)
        assert res.status_code == 200
Lianmin Zheng's avatar
Lianmin Zheng committed
228
        scores.append(res.json().get("prompt_score", 0))
Lianmin Zheng's avatar
Lianmin Zheng committed
229
230
231
232
233
234
235
236
237
238
239
    return np.argmax(scores)

    """
    Modify vllm/entrypoints/api_server.py

    if final_output.prompt_logprobs is not None:
        score = np.mean([prob[t_id] for t_id, prob in zip(final_output.prompt_token_ids[1:], final_output.prompt_logprobs[1:])])
        ret["prompt_score"] = score
    """


Liangsheng Yin's avatar
Liangsheng Yin committed
240
241
242
243
244
245
246
247
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"])


248
def add_common_other_args_and_parse(parser: argparse.ArgumentParser):
Lianmin Zheng's avatar
Lianmin Zheng committed
249
    parser.add_argument("--parallel", type=int, default=64)
Lianmin Zheng's avatar
Lianmin Zheng committed
250
251
252
253
254
255
    parser.add_argument("--host", type=str, default="http://127.0.0.1")
    parser.add_argument("--port", type=int, default=None)
    parser.add_argument(
        "--backend",
        type=str,
        required=True,
Liangsheng Yin's avatar
Liangsheng Yin committed
256
257
258
259
        choices=[
            "vllm",
            "outlines",
            "lightllm",
260
            "gserver",
Liangsheng Yin's avatar
Liangsheng Yin committed
261
262
263
264
            "guidance",
            "srt-raw",
            "llama.cpp",
        ],
Lianmin Zheng's avatar
Lianmin Zheng committed
265
    )
Liangsheng Yin's avatar
Liangsheng Yin committed
266
    parser.add_argument("--n-ctx", type=int, default=4096)
Lianmin Zheng's avatar
Lianmin Zheng committed
267
268
269
270
271
272
273
274
275
    parser.add_argument(
        "--model-path", type=str, default="meta-llama/Llama-2-7b-chat-hf"
    )
    parser.add_argument("--result-file", type=str, default="result.jsonl")
    args = parser.parse_args()

    if args.port is None:
        default_port = {
            "vllm": 21000,
Liangsheng Yin's avatar
Liangsheng Yin committed
276
            "outlines": 21000,
Lianmin Zheng's avatar
Lianmin Zheng committed
277
278
            "lightllm": 22000,
            "srt-raw": 30000,
279
            "gserver": 9988,
Lianmin Zheng's avatar
Lianmin Zheng committed
280
281
282
283
284
        }
        args.port = default_port.get(args.backend, None)
    return args


285
def add_common_sglang_args_and_parse(parser: argparse.ArgumentParser):
Lianmin Zheng's avatar
Lianmin Zheng committed
286
287
288
289
290
291
292
293
294
    parser.add_argument("--parallel", type=int, default=64)
    parser.add_argument("--host", type=str, default="http://127.0.0.1")
    parser.add_argument("--port", type=int, default=30000)
    parser.add_argument("--backend", type=str, default="srt")
    parser.add_argument("--result-file", type=str, default="result.jsonl")
    args = parser.parse_args()
    return args


295
def select_sglang_backend(args: argparse.Namespace):
Lianmin Zheng's avatar
Lianmin Zheng committed
296
297
298
299
    if args.backend.startswith("srt"):
        if args.backend == "srt-no-parallel":
            global_config.enable_parallel_encoding = False
        backend = RuntimeEndpoint(f"{args.host}:{args.port}")
300
    elif args.backend.startswith("gpt-"):
Lianmin Zheng's avatar
Lianmin Zheng committed
301
302
303
304
        backend = OpenAI(args.backend)
    else:
        raise ValueError(f"Invalid backend: {args.backend}")
    return backend
Liangsheng Yin's avatar
Liangsheng Yin committed
305
306


307
def _get_call_generate(args: argparse.Namespace):
Liangsheng Yin's avatar
Liangsheng Yin committed
308
309
310
311
312
313
    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")
314
315
    elif args.backend == "gserver":
        return partial(call_generate_gserver, url=f"{args.host}:{args.port}")
Liangsheng Yin's avatar
Liangsheng Yin committed
316
317
318
319
320
321
322
323
324
325
326
327
328
    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
    else:
        raise ValueError(f"Invalid backend: {args.backend}")


329
def _get_call_select(args: argparse.Namespace):
Liangsheng Yin's avatar
Liangsheng Yin committed
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
    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
    else:
        raise ValueError(f"Invalid backend: {args.backend}")


346
def get_call_generate(args: argparse.Namespace):
Liangsheng Yin's avatar
Liangsheng Yin committed
347
348
349
350
351
352
353
354
355
356
357
358
    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


359
def get_call_select(args: argparse.Namespace):
Liangsheng Yin's avatar
Liangsheng Yin committed
360
361
362
363
364
365
366
367
368
369
    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
370
371


372
def popen_launch_server(
373
374
375
376
    model: str,
    base_url: str,
    timeout: float,
    api_key: Optional[str] = None,
Mick's avatar
Mick committed
377
    other_args: list[str] = (),
378
    env: Optional[dict] = None,
379
    return_stdout_stderr: Optional[tuple] = None,
380
    pd_seperated: bool = False,
381
382
383
384
):
    _, host, port = base_url.split(":")
    host = host[2:]

385
386
387
388
389
    if pd_seperated:
        command = "sglang.launch_pd_server"
    else:
        command = "sglang.launch_server"

390
391
392
    command = [
        "python3",
        "-m",
393
        command,
394
395
        "--model-path",
        model,
396
        *[str(x) for x in other_args],
397
    ]
Chayenne's avatar
Chayenne committed
398

399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
    if pd_seperated:
        command.extend(
            [
                "--lb-host",
                host,
                "--lb-port",
                port,
            ]
        )
    else:
        command.extend(
            [
                "--host",
                host,
                "--port",
                port,
            ]
        )

418
419
420
    if api_key:
        command += ["--api-key", api_key]

421
422
    print(f"command={' '.join(command)}")

423
424
425
    if return_stdout_stderr:
        process = subprocess.Popen(
            command,
426
427
            stdout=return_stdout_stderr[0],
            stderr=return_stdout_stderr[1],
428
429
430
431
432
            env=env,
            text=True,
        )
    else:
        process = subprocess.Popen(command, stdout=None, stderr=None, env=env)
433
434

    start_time = time.time()
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
    with requests.Session() as session:
        while time.time() - start_time < timeout:
            try:
                headers = {
                    "Content-Type": "application/json; charset=utf-8",
                    "Authorization": f"Bearer {api_key}",
                }
                response = session.get(
                    f"{base_url}/health_generate",
                    headers=headers,
                )
                if response.status_code == 200:
                    return process
            except requests.RequestException:
                pass
450
451
452
453
454

            return_code = process.poll()
            if return_code is not None:
                raise Exception(f"Server unexpectedly exits ({return_code=}).")

455
            time.sleep(10)
456
457

    kill_process_tree(process.pid)
458
    raise TimeoutError("Server failed to start within the timeout period.")
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484


def run_with_timeout(
    func: Callable,
    args: tuple = (),
    kwargs: Optional[dict] = None,
    timeout: float = None,
):
    """Run a function with timeout."""
    ret_value = []

    def _target_func():
        ret_value.append(func(*args, **(kwargs or {})))

    t = threading.Thread(target=_target_func)
    t.start()
    t.join(timeout=timeout)
    if t.is_alive():
        raise TimeoutError()

    if not ret_value:
        raise RuntimeError()

    return ret_value[0]


Byron Hsu's avatar
Byron Hsu committed
485
486
487
488
489
490
491
@dataclass
class TestFile:
    name: str
    estimated_time: float = 60


def run_unittest_files(files: List[TestFile], timeout_per_file: float):
492
493
494
    tic = time.time()
    success = True

Lianmin Zheng's avatar
Lianmin Zheng committed
495
496
    for file in files:
        filename, estimated_time = file.name, file.estimated_time
497
        process = None
498

Mingyi's avatar
Mingyi committed
499
        def run_one_file(filename):
500
501
            nonlocal process

Mingyi's avatar
Mingyi committed
502
            filename = os.path.join(os.getcwd(), filename)
Lianmin Zheng's avatar
Lianmin Zheng committed
503
504
505
            print(f".\n.\nBegin:\npython3 {filename}\n.\n.\n", flush=True)
            tic = time.time()

Mingyi's avatar
Mingyi committed
506
507
508
509
            process = subprocess.Popen(
                ["python3", filename], stdout=None, stderr=None, env=os.environ
            )
            process.wait()
Lianmin Zheng's avatar
Lianmin Zheng committed
510
511
512
513
514
515
            elapsed = time.time() - tic

            print(
                f".\n.\nEnd:\n{filename=}, {elapsed=:.0f}, {estimated_time=}\n.\n.\n",
                flush=True,
            )
Mingyi's avatar
Mingyi committed
516
            return process.returncode
517
518

        try:
Mingyi's avatar
Mingyi committed
519
520
521
            ret_code = run_with_timeout(
                run_one_file, args=(filename,), timeout=timeout_per_file
            )
522
523
524
            assert (
                ret_code == 0
            ), f"expected return code 0, but {filename} returned {ret_code}"
525
        except TimeoutError:
526
            kill_process_tree(process.pid)
527
528
            time.sleep(5)
            print(
529
530
                f"\nTimeout after {timeout_per_file} seconds when running {filename}\n",
                flush=True,
531
            )
Mingyi's avatar
Mingyi committed
532
533
            success = False
            break
534
535

    if success:
536
        print(f"Success. Time elapsed: {time.time() - tic:.2f}s", flush=True)
537
    else:
538
        print(f"Fail. Time elapsed: {time.time() - tic:.2f}s", flush=True)
539
540

    return 0 if success else -1
541
542
543
544


def get_similarities(vec1, vec2):
    return F.cosine_similarity(torch.tensor(vec1), torch.tensor(vec2), dim=0)
545
546


547
548
549
550
551
552
def get_benchmark_args(
    base_url="",
    dataset_name="",
    dataset_path="",
    tokenizer="",
    num_prompts=500,
553
    sharegpt_output_len=None,
554
555
    random_input_len=4096,
    random_output_len=2048,
556
    sharegpt_context_len=None,
557
558
559
    request_rate=float("inf"),
    disable_stream=False,
    disable_ignore_eos=False,
560
    seed: int = 0,
561
    pd_seperated: bool = False,
562
563
564
565
566
567
568
569
570
571
572
):
    return SimpleNamespace(
        backend="sglang",
        base_url=base_url,
        host=None,
        port=None,
        dataset_name=dataset_name,
        dataset_path=dataset_path,
        model=None,
        tokenizer=tokenizer,
        num_prompts=num_prompts,
573
574
        sharegpt_output_len=sharegpt_output_len,
        sharegpt_context_len=sharegpt_context_len,
575
576
577
578
579
580
581
582
583
        random_input_len=random_input_len,
        random_output_len=random_output_len,
        random_range_ratio=0.0,
        request_rate=request_rate,
        multi=None,
        output_file=None,
        disable_tqdm=False,
        disable_stream=disable_stream,
        return_logprob=False,
584
        seed=seed,
585
586
587
588
589
        disable_ignore_eos=disable_ignore_eos,
        extra_request_body=None,
        apply_chat_template=False,
        profile=None,
        lora_name=None,
590
591
        prompt_suffix="",
        pd_seperated=pd_seperated,
592
593
594
    )


595
596
597
598
599
600
def run_bench_serving(
    model,
    num_prompts,
    request_rate,
    other_server_args,
    dataset_name="random",
601
602
    dataset_path="",
    tokenizer=None,
603
604
    random_input_len=4096,
    random_output_len=2048,
605
    sharegpt_context_len=None,
606
    disable_stream=False,
607
    disable_ignore_eos=False,
608
    need_warmup=False,
609
    seed: int = 0,
610
):
611
612
613
614
615
616
617
618
619
620
    # 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
621
    args = get_benchmark_args(
622
        base_url=base_url,
623
        dataset_name=dataset_name,
624
625
        dataset_path=dataset_path,
        tokenizer=tokenizer,
626
        num_prompts=num_prompts,
627
628
        random_input_len=random_input_len,
        random_output_len=random_output_len,
629
        sharegpt_context_len=sharegpt_context_len,
630
        request_rate=request_rate,
631
        disable_stream=disable_stream,
632
        disable_ignore_eos=disable_ignore_eos,
633
        seed=seed,
634
635
636
    )

    try:
637
638
639
640
        if need_warmup:
            warmup_args = copy.deepcopy(args)
            warmup_args.num_prompts = 16
            run_benchmark(warmup_args)
641
642
        res = run_benchmark(args)
    finally:
643
        kill_process_tree(process.pid)
644
645
646

    assert res["completed"] == num_prompts
    return res
647
648


649
650
651
652
653
654
def run_bench_serving_multi(
    model,
    base_url,
    other_server_args,
    benchmark_args,
    need_warmup=False,
655
    pd_seperated=False,
656
657
658
659
660
661
662
):
    # Launch the server
    process = popen_launch_server(
        model,
        base_url,
        timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
        other_args=other_server_args,
663
        pd_seperated=pd_seperated,
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
    )

    # run benchmark for all
    res_l = []
    try:
        for args in benchmark_args:
            if need_warmup:
                warmup_args = copy.deepcopy(args)
                warmup_args.num_prompts = 16
                run_benchmark(warmup_args)

            res = run_benchmark(args)
            res_l.append((args, res))
    finally:
        kill_process_tree(process.pid)

    return res_l


683
def run_bench_one_batch(model, other_args):
684
685
686
    command = [
        "python3",
        "-m",
687
        "sglang.bench_one_batch",
688
689
690
691
692
693
        "--batch-size",
        "1",
        "--input",
        "128",
        "--output",
        "8",
694
        *[str(x) for x in other_args],
695
    ]
saienduri's avatar
saienduri committed
696
697
    if model is not None:
        command += ["--model-path", model]
698
699
700
701
702
703
704
705
706
707
708
709
    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:
710
        kill_process_tree(process.pid)
711
712

    return output_throughput
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746


def lcs(X, Y):
    m = len(X)
    n = len(Y)
    L = [[0] * (n + 1) for _ in range(m + 1)]

    for i in range(m + 1):
        for j in range(n + 1):
            if i == 0 or j == 0:
                L[i][j] = 0
            elif X[i - 1] == Y[j - 1]:
                L[i][j] = L[i - 1][j - 1] + 1
            else:
                L[i][j] = max(L[i - 1][j], L[i][j - 1])

    return L[m][n]


def calculate_rouge_l(output_strs_list1, output_strs_list2):
    """calculate the ROUGE-L score"""
    rouge_l_scores = []

    for s1, s2 in zip(output_strs_list1, output_strs_list2):
        lcs_len = lcs(s1, s2)
        precision = lcs_len / len(s1) if len(s1) > 0 else 0
        recall = lcs_len / len(s2) if len(s2) > 0 else 0
        if precision + recall > 0:
            fmeasure = (2 * precision * recall) / (precision + recall)
        else:
            fmeasure = 0.0
        rouge_l_scores.append(fmeasure)

    return rouge_l_scores
747
748
749


STDERR_FILENAME = "stderr.txt"
750
STDOUT_FILENAME = "stdout.txt"
751
752


753
def read_output(output_lines: List[str], filename: str = STDERR_FILENAME):
754
    """Print the output in real time with another thread."""
755
    while not os.path.exists(filename):
756
757
        time.sleep(1)

758
759
    pt = 0
    while pt >= 0:
760
        if pt > 0 and not os.path.exists(filename):
761
            break
762
        lines = open(filename).readlines()
763
764
        for line in lines[pt:]:
            print(line, end="", flush=True)
765
            output_lines.append(line)
766
            pt += 1
767
        time.sleep(0.1)
768
769


770
771
def run_and_check_memory_leak(
    workload_func,
772
    disable_radix_cache,
773
    enable_mixed_chunk,
774
    disable_overlap,
775
    chunked_prefill_size,
776
    assert_has_abort,
777
):
778
779
780
781
782
783
    other_args = [
        "--chunked-prefill-size",
        str(chunked_prefill_size),
        "--log-level",
        "debug",
    ]
784
785
786
787
    if disable_radix_cache:
        other_args += ["--disable-radix-cache"]
    if enable_mixed_chunk:
        other_args += ["--enable-mixed-chunk"]
788
789
    if disable_overlap:
        other_args += ["--disable-overlap-schedule"]
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810

    model = DEFAULT_MODEL_NAME_FOR_TEST
    port = random.randint(4000, 5000)
    base_url = f"http://127.0.0.1:{port}"

    # Create files and launch the server
    stdout = open(STDOUT_FILENAME, "w")
    stderr = open(STDERR_FILENAME, "w")
    process = popen_launch_server(
        model,
        base_url,
        timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
        other_args=other_args,
        return_stdout_stderr=(stdout, stderr),
    )

    # Launch a thread to stream the output
    output_lines = []
    t = threading.Thread(target=read_output, args=(output_lines,))
    t.start()

811
812
    # Run the workload
    workload_func(base_url, model)
813
814

    # Clean up everything
815
    kill_process_tree(process.pid)
816
817
    stdout.close()
    stderr.close()
818
819
820
821
    if os.path.exists(STDOUT_FILENAME):
        os.remove(STDOUT_FILENAME)
    if os.path.exists(STDERR_FILENAME):
        os.remove(STDERR_FILENAME)
Lianmin Zheng's avatar
Lianmin Zheng committed
822
    kill_process_tree(process.pid)
823
824
825
826
827
    t.join()

    # Assert success
    has_new_server = False
    has_leak = False
828
    has_abort = False
829
    for line in output_lines:
Lianmin Zheng's avatar
Lianmin Zheng committed
830
        if "Uvicorn running" in line:
831
832
833
            has_new_server = True
        if "leak" in line:
            has_leak = True
834
835
        if "Abort" in line:
            has_abort = True
836
837

    assert has_new_server
838
    assert not has_leak
839
840
    if assert_has_abort:
        assert has_abort
841
842


843
844
845
846
def run_command_and_capture_output(command, env: Optional[dict] = None):
    stdout = open(STDOUT_FILENAME, "w")
    stderr = open(STDERR_FILENAME, "w")
    process = subprocess.Popen(
847
        command, stdout=stdout, stderr=stdout, env=env, text=True
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
    )

    # Launch a thread to stream the output
    output_lines = []
    t = threading.Thread(target=read_output, args=(output_lines, STDOUT_FILENAME))
    t.start()

    # Join the process
    process.wait()

    stdout.close()
    stderr.close()
    if os.path.exists(STDOUT_FILENAME):
        os.remove(STDOUT_FILENAME)
    if os.path.exists(STDERR_FILENAME):
        os.remove(STDERR_FILENAME)
    kill_process_tree(process.pid)
    t.join()

    return output_lines


870
871
872
def run_mmlu_test(
    disable_radix_cache=False,
    enable_mixed_chunk=False,
873
    disable_overlap=False,
874
875
876
877
878
879
880
881
882
883
884
885
886
887
    chunked_prefill_size=32,
):
    def workload_func(base_url, model):
        # Run the eval
        args = SimpleNamespace(
            base_url=base_url,
            model=model,
            eval_name="mmlu",
            num_examples=128,
            num_threads=128,
        )

        try:
            metrics = run_eval(args)
Lianmin Zheng's avatar
Lianmin Zheng committed
888
            assert metrics["score"] >= 0.65, f"{metrics=}"
889
890
891
        finally:
            pass

Chayenne's avatar
Chayenne committed
892
893
894
895
    run_and_check_memory_leak(
        workload_func,
        disable_radix_cache,
        enable_mixed_chunk,
896
        disable_overlap,
Chayenne's avatar
Chayenne committed
897
        chunked_prefill_size,
898
        assert_has_abort=False,
Chayenne's avatar
Chayenne committed
899
    )
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930


def run_mulit_request_test(
    disable_radix_cache=False,
    enable_mixed_chunk=False,
    enable_overlap=False,
    chunked_prefill_size=32,
):
    def workload_func(base_url, model):
        def run_one(_):
            prompt = """
            System: You are a helpful assistant.
            User: What is the capital of France?
            Assistant: The capital of France is
            """

            response = requests.post(
                f"{base_url}/generate",
                json={
                    "text": prompt,
                    "sampling_params": {
                        "temperature": 0,
                        "max_new_tokens": 8,
                    },
                },
            )
            ret = response.json()

        with ThreadPoolExecutor(2) as executor:
            list(executor.map(run_one, list(range(4))))

Chayenne's avatar
Chayenne committed
931
932
933
934
935
936
    run_and_check_memory_leak(
        workload_func,
        disable_radix_cache,
        enable_mixed_chunk,
        enable_overlap,
        chunked_prefill_size,
937
        assert_has_abort=False,
Chayenne's avatar
Chayenne committed
938
    )
939
940
941


def write_github_step_summary(content):
942
943
944
945
    if not os.environ.get("GITHUB_STEP_SUMMARY"):
        logging.warning("GITHUB_STEP_SUMMARY environment variable not set")
        return

946
947
    with open(os.environ["GITHUB_STEP_SUMMARY"], "a") as f:
        f.write(content)
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022


def run_logprob_check(self: unittest.TestCase, arg: Tuple):
    (
        input_len,
        output_len,
        temperature,
        logprob_start_len,
        return_logprob,
        top_logprobs_num,
    ) = arg
    input_ids = list(range(input_len))

    response = requests.post(
        self.base_url + "/generate",
        json={
            "input_ids": input_ids,
            "sampling_params": {
                "temperature": temperature,
                "max_new_tokens": output_len,
                "ignore_eos": True,
            },
            "return_logprob": return_logprob,
            "logprob_start_len": logprob_start_len,
            "top_logprobs_num": top_logprobs_num,
        },
    )
    response_json = response.json()

    res = response_json
    self.assertEqual(res["meta_info"]["prompt_tokens"], input_len)
    self.assertEqual(res["meta_info"]["completion_tokens"], output_len)

    # Test the number of tokens are correct
    if return_logprob:
        self.assertEqual(
            len(res["meta_info"]["input_token_logprobs"]) + logprob_start_len,
            res["meta_info"]["prompt_tokens"],
        )
        self.assertEqual(len(res["meta_info"]["output_token_logprobs"]), output_len)

        if top_logprobs_num:
            self.assertEqual(
                len(res["meta_info"]["input_top_logprobs"]) + logprob_start_len,
                res["meta_info"]["prompt_tokens"],
            )
            self.assertEqual(len(res["meta_info"]["output_top_logprobs"]), output_len)

            for i in range(output_len):
                self.assertEqual(
                    len(res["meta_info"]["output_top_logprobs"][i]),
                    top_logprobs_num,
                )

                # Test the top-1 tokens are the same as output tokens if temperature == 0
                if temperature == 0:
                    rank = 0
                    while rank < len(res["meta_info"]["output_top_logprobs"][i]):
                        try:
                            self.assertListEqual(
                                res["meta_info"]["output_token_logprobs"][i],
                                res["meta_info"]["output_top_logprobs"][i][rank],
                            )
                            break
                        except AssertionError:
                            # There's a tie. Allow the second item in this case.
                            if (
                                res["meta_info"]["output_top_logprobs"][i][rank][0]
                                == res["meta_info"]["output_top_logprobs"][i][rank + 1][
                                    0
                                ]
                            ):
                                rank += 1
                            else:
                                raise
1023
1024
1025
1026


class CustomTestCase(unittest.TestCase):
    def _callTestMethod(self, method):
1027
        max_retry = int(
Yineng Zhang's avatar
Yineng Zhang committed
1028
            os.environ.get("SGLANG_TEST_MAX_RETRY", "1" if is_in_ci() else "0")
1029
        )
1030
1031
1032
        retry(
            lambda: super(CustomTestCase, self)._callTestMethod(method),
            max_retry=max_retry,
1033
        )