MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT-LLM / main

Function main

tensorrt_llm/serve/scripts/benchmark_serving.py:689–1015  ·  view source on GitHub ↗
(args: argparse.Namespace)

Source from the content-addressed store, hash-verified

687
688
689def main(args: argparse.Namespace):
690 print(args)
691 random.seed(args.seed)
692 np.random.seed(args.seed)
693
694 backend = args.backend
695 model_id = args.model
696 model_name = args.served_model_name
697 tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model
698 tokenizer_mode = args.tokenizer_mode
699
700 if backend == "openai-chat":
701 args.endpoint = "/v1/chat/completions"
702
703 if args.base_url is not None:
704 api_url = f"{args.base_url}{args.endpoint}"
705 base_url = f"{args.base_url}"
706 else:
707 api_url = f"http://{args.host}:{args.port}{args.endpoint}"
708 base_url = f"http://{args.host}:{args.port}"
709
710 tokenizer = get_tokenizer(tokenizer_id,
711 tokenizer_mode=tokenizer_mode,
712 trust_remote_code=args.trust_remote_code)
713
714 if args.dataset_name is None:
715 raise ValueError(
716 "Please specify '--dataset-name' and the corresponding "
717 "'--dataset-path' if required.")
718
719 if args.dataset_name == "sonnet":
720 dataset = SonnetDataset(dataset_path=args.dataset_path)
721 # For the "sonnet" dataset, formatting depends on the backend.
722 if args.backend == "openai-chat":
723 input_requests = dataset.sample(num_requests=args.num_prompts,
724 input_len=args.sonnet_input_len,
725 output_len=args.sonnet_output_len,
726 prefix_len=args.sonnet_prefix_len,
727 tokenizer=tokenizer,
728 return_prompt_formatted=False)
729 else:
730 assert tokenizer.chat_template or tokenizer.default_chat_template, (
731 "Tokenizer/model must have chat template for sonnet dataset.")
732 input_requests = dataset.sample(num_requests=args.num_prompts,
733 input_len=args.sonnet_input_len,
734 output_len=args.sonnet_output_len,
735 prefix_len=args.sonnet_prefix_len,
736 tokenizer=tokenizer,
737 return_prompt_formatted=True)
738
739 elif args.dataset_name == "hf":
740 # all following datasets are implemented from the
741 # HuggingFaceDataset base class
742 if args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS:
743 dataset_class = VisionArenaDataset
744 args.hf_split = "train"
745 args.hf_subset = None
746 elif args.dataset_path in InstructCoderDataset.SUPPORTED_DATASET_PATHS:

Callers 1

Calls 15

sampleMethod · 0.95
parse_json_fileMethod · 0.95
create_timing_diagramMethod · 0.95
get_tokenizerFunction · 0.90
SonnetDatasetClass · 0.90
CustomDatasetClass · 0.90
check_goodput_argsFunction · 0.85
fetch_perf_metricsFunction · 0.85
benchmarkFunction · 0.70

Tested by

no test coverage detected