83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
import argparse
|
|
|
|
from sglang.auto_benchmark_lib import (
|
|
SUPPORTED_DATASETS,
|
|
convert_dataset,
|
|
run_auto_benchmark,
|
|
validate_dataset,
|
|
)
|
|
|
|
|
|
def add_dataset_args(parser: argparse.ArgumentParser) -> None:
|
|
parser.add_argument(
|
|
"--kind",
|
|
required=True,
|
|
choices=sorted(SUPPORTED_DATASETS),
|
|
help="Dataset kind: sharegpt, custom, random, or generated-shared-prefix.",
|
|
)
|
|
parser.add_argument(
|
|
"--path",
|
|
default="",
|
|
help="Dataset file path. Leave empty for sharegpt auto-download.",
|
|
)
|
|
parser.add_argument("--tokenizer", required=True)
|
|
parser.add_argument("--model", default=None)
|
|
parser.add_argument("--num-prompts", type=int, default=1000)
|
|
parser.add_argument("--output-len", type=int, default=None)
|
|
parser.add_argument("--context-len", type=int, default=None)
|
|
parser.add_argument("--prompt-suffix", type=str, default="")
|
|
parser.add_argument("--apply-chat-template", action="store_true")
|
|
parser.add_argument("--random-input-len", type=int, default=1024)
|
|
parser.add_argument("--random-output-len", type=int, default=256)
|
|
parser.add_argument("--random-range-ratio", type=float, default=0.0)
|
|
parser.add_argument("--gsp-num-groups", type=int, default=64)
|
|
parser.add_argument("--gsp-prompts-per-group", type=int, default=16)
|
|
parser.add_argument("--gsp-system-prompt-len", type=int, default=2048)
|
|
parser.add_argument("--gsp-question-len", type=int, default=128)
|
|
parser.add_argument("--gsp-output-len", type=int, default=256)
|
|
parser.add_argument("--gsp-range-ratio", type=float, default=1.0)
|
|
parser.add_argument("--gsp-fast-prepare", action="store_true")
|
|
parser.add_argument("--gsp-send-routing-key", action="store_true")
|
|
parser.add_argument("--gsp-num-turns", type=int, default=1)
|
|
parser.add_argument("--gsp-ordered", action="store_true")
|
|
parser.add_argument("--seed", type=int, default=1)
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description="SGLang auto benchmark utilities.")
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
run_parser = subparsers.add_parser(
|
|
"run", help="Run auto benchmark from YAML config."
|
|
)
|
|
run_parser.add_argument("--config", required=True)
|
|
|
|
convert_parser = subparsers.add_parser(
|
|
"convert",
|
|
help="Prepare sharegpt/custom/random/generated-shared-prefix data into canonical autobench JSONL.",
|
|
)
|
|
add_dataset_args(convert_parser)
|
|
convert_parser.add_argument("--output", required=True)
|
|
|
|
validate_parser = subparsers.add_parser(
|
|
"validate", help="Validate a canonical autobench JSONL dataset."
|
|
)
|
|
validate_parser.add_argument("--dataset-path", required=True)
|
|
validate_parser.add_argument("--tokenizer", required=True)
|
|
|
|
return parser
|
|
|
|
|
|
def main() -> None:
|
|
args = build_parser().parse_args()
|
|
if args.command == "run":
|
|
run_auto_benchmark(args.config)
|
|
elif args.command == "convert":
|
|
convert_dataset(args)
|
|
elif args.command == "validate":
|
|
validate_dataset(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|