Function main

benchmarks/benchmark_serving.py:1043–1222

Main entry point

(args: argparse.Namespace)

def main(args: argparse.Namespace):
    """Main entry point"""
    print(args)
    random.seed(args.seed)
    np.random.seed(args.seed)

    backend = args.backend
    # 支持多轮对话方式请求，仅支持chat接口
    if args.multi_turn:
        backend = "openai-chat-multi-turn"
    model_id = args.model
    model_name = args.served_model_name
    tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model

    if args.base_url is not None:
        api_url = f"{args.base_url}{args.endpoint}"
        base_url = f"{args.base_url}"
    else:
        api_url = f"http://{args.host}:{args.port}{args.endpoint}"
        base_url = f"http://{args.host}:{args.port}"

    if args.dataset_name is None:
        raise ValueError("Please specify '--dataset-name' and the corresponding " "'--dataset-path' if required.")

    # For datasets that follow a similar structure, use a mapping.
    dataset_mapping = {
        "EB": lambda: EBDataset(random_seed=args.seed, dataset_path=args.dataset_path, shuffle=args.shuffle).sample(
            num_requests=args.num_prompts,
            output_len=args.sharegpt_output_len,
        ),
        "EBChat": lambda: EBChatDataset(
            random_seed=args.seed, dataset_path=args.dataset_path, shuffle=args.shuffle
        ).sample(
            num_requests=args.num_prompts,
            output_len=args.sharegpt_output_len,
        ),
        "random": lambda: RandomTextDataset().sample(
            num_requests=args.num_prompts,
            random_input_len=args.random_input_len,
            random_output_len=args.random_output_len,
            random_range_ratio=args.random_range_ratio,
        ),
    }

    try:
        input_requests = dataset_mapping[args.dataset_name]()
    except KeyError as err:
        raise ValueError(f"Unknown dataset: {args.dataset_name}") from err

    goodput_config_dict = check_goodput_args(args)

    # Collect the sampling parameters.
    sampling_params = {
        k: v
        for k, v in {
            "top_p": args.top_p,
            "top_k": args.top_k,
            "min_p": args.min_p,