wehub-resource-sync
430 lines
16 KiB
Python
430 lines
16 KiB
Python
import json
|
|
import multiprocessing as mp
|
|
import os
|
|
import random
|
|
import sys
|
|
import time
|
|
from typing import List, Tuple
|
|
import numpy as np
|
|
import requests
|
|
from sentencepiece import SentencePieceProcessor
|
|
from typing import AsyncGenerator, List, Optional, Tuple, Union
|
|
import argparse
|
|
|
|
|
|
def infer(
|
|
session_id: int,
|
|
req_que: mp.Queue,
|
|
res_que: mp.Queue,
|
|
end_flags: mp.Queue,
|
|
backend: str = "tgi",
|
|
api_url: str = "http://localhost:8010/generate_stream",
|
|
model_name: str = "openlm-research/open_llama_13b",
|
|
min_dec_len: int = 1,
|
|
max_dec_len: int = 2048
|
|
):
|
|
stats = []
|
|
while not req_que.empty():
|
|
try:
|
|
prompt, input_seqlen, output_seqlen = req_que.get(timeout=10.0)
|
|
except:
|
|
continue
|
|
|
|
start = time.time()
|
|
is_first = True
|
|
first_token_latency = float("inf")
|
|
# if output_seqlen > 1024:
|
|
# print("Request exceeds 1024 tokens. Truncating to 1024.", output_seqlen)
|
|
# output_seqlen = 1024
|
|
|
|
headers = {"User-Agent": "Benchmark Client"}
|
|
if backend == "vllm":
|
|
pload = {
|
|
"model": model_name,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0.6,
|
|
"top_p": 0.95,
|
|
"max_tokens": output_seqlen,
|
|
"stream": True
|
|
}
|
|
elif backend == "sglang":
|
|
pload = {
|
|
"text": prompt,
|
|
"sampling_params": {
|
|
"max_new_tokens": max_dec_len,
|
|
"temperature": 0.6,
|
|
"top_k": 1,
|
|
"top_p": 0.95,
|
|
},
|
|
"return_logprob": False,
|
|
"stream": True
|
|
}
|
|
elif backend == "paddle":
|
|
pload = {
|
|
"text": prompt,
|
|
"max_dec_len": max_dec_len,
|
|
"min_dec_len": min_dec_len,
|
|
"topp": 0.95,
|
|
"temperature": 0.6,
|
|
"stream": True,
|
|
"return_all_tokens": False
|
|
}
|
|
else: # backend is trtllm
|
|
pload = {
|
|
"model": model_name,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0.6,
|
|
"top_p": 0.95,
|
|
"max_tokens": max_dec_len,
|
|
"stream": True
|
|
}
|
|
|
|
response = requests.post(url=api_url, headers=headers, json=pload, stream=True)
|
|
|
|
chunks = []
|
|
for chunk in response.iter_content(chunk_size=1000000):
|
|
chunks.append(chunk)
|
|
if is_first:
|
|
first_token_latency = time.time() - start
|
|
is_first = False
|
|
|
|
total_token_latency = time.time() - start
|
|
inference_first_token_latency = float(
|
|
json.loads(chunks[0])["inference_time_cost"]
|
|
)
|
|
inference_total_token_latency = float(
|
|
json.loads(chunks[-1])["inference_time_cost"]
|
|
)
|
|
|
|
if backend == "vllm":
|
|
try:
|
|
res_text = b"".join(chunks).decode("utf-8")
|
|
token_num = eval(
|
|
res_text.split("previous_num_tokens:")[-1].split("data: [DONE]")[0]
|
|
)
|
|
except:
|
|
token_num = len(chunks)
|
|
elif backend == "paddle":
|
|
token_num = 0
|
|
for chunk in chunks:
|
|
chunk_dict = json.loads(chunk.decode("utf-8").strip())
|
|
if chunk_dict["is_end"]:
|
|
token_num += chunk_dict["tokens_all_num"]
|
|
else:
|
|
token_num = len(chunks)
|
|
|
|
stats.append(
|
|
[first_token_latency, total_token_latency, input_seqlen, output_seqlen, token_num, inference_first_token_latency, inference_total_token_latency]
|
|
)
|
|
print(f"Request queue size: {req_que.qsize()}, Real return tokens: {token_num}, Request Chunks: {len(chunks)}, Label out_seq_len: {output_seqlen}")
|
|
|
|
print(f"Process ID {os.getpid()} has processed all requests.")
|
|
if len(stats) > 0:
|
|
res_que.put((session_id, stats))
|
|
end_flags.put(1)
|
|
print(f"Session {session_id} (PID={os.getpid()}) completed.")
|
|
|
|
|
|
def warmup(
|
|
concurrency: int,
|
|
output_seqlen: int,
|
|
warmup_round: int = 4,
|
|
backend: str = "tgi",
|
|
api_url: str = "http://localhost:8010/generate_stream",
|
|
model_name: str = "openlm-research/open_llama_13b",
|
|
):
|
|
print("Starting warmup process...")
|
|
|
|
def _infer(index, warmup_round, backend):
|
|
headers = {"User-Agent": "Benchmark Client"}
|
|
if backend == "vllm":
|
|
pload = {
|
|
"model": model_name,
|
|
"messages": [{"role": "user", "content": "What is the capital of France?"}],
|
|
"temperature": 1.0,
|
|
"top_p": 0.0,
|
|
"stream": True,
|
|
}
|
|
elif backend == "sglang":
|
|
pload = {
|
|
"text": "how about beijing?",
|
|
"sampling_params": {
|
|
"max_new_tokens": output_seqlen,
|
|
"temperature": 1.0,
|
|
"top_k": 1,
|
|
"top_p": 0.0,
|
|
},
|
|
"return_logprob": False,
|
|
"stream": True
|
|
}
|
|
elif backend == "paddle":
|
|
pload = {
|
|
"text": "how about beijing?",
|
|
"max_dec_len": output_seqlen,
|
|
"topp": 1.0,
|
|
"temperature": 1.0,
|
|
"stream": True,
|
|
"return_all_tokens": False
|
|
}
|
|
else: # trtllm
|
|
pload = {
|
|
"model": model_name,
|
|
"messages": [{"role": "user", "content": "What is the capital of France?"}],
|
|
"temperature": 1.0,
|
|
"top_p": 1.0,
|
|
"max_tokens": 50,
|
|
"stream": True
|
|
}
|
|
|
|
for _ in range(warmup_round):
|
|
response = requests.post(url=api_url, headers=headers, json=pload, stream=True)
|
|
for _ in response.iter_content(chunk_size=1024):
|
|
pass
|
|
|
|
start_time = time.perf_counter()
|
|
procs = []
|
|
for i in range(concurrency):
|
|
proc = mp.Process(target=_infer, args=(i, warmup_round, backend))
|
|
procs.append(proc)
|
|
proc.start()
|
|
for proc in procs:
|
|
proc.join()
|
|
end_time = time.perf_counter()
|
|
print(f"Warmup completed. Elapsed time: {round(end_time - start_time, 2)} seconds")
|
|
|
|
|
|
def sample_requests_filtered_shared_gpt(dataset_path: str, num_requests: int):
|
|
with open(dataset_path, "r") as infile:
|
|
filtered_dataset = json.load(infile)
|
|
extracted_data = [[entry[0], entry[1], entry[2]] for entry in filtered_dataset[:num_requests]]
|
|
que = mp.Queue()
|
|
|
|
total_input_len = sum(entry[1] for entry in extracted_data)
|
|
total_output_len = sum(entry[2] for entry in extracted_data)
|
|
avg_input_len = total_input_len / len(extracted_data)
|
|
avg_output_len = total_output_len / len(extracted_data)
|
|
|
|
print(f"Number of test samples: {len(extracted_data)}")
|
|
print(f"Average input length: {avg_input_len:.2f}")
|
|
print(f"Average output length: {avg_output_len:.2f}")
|
|
|
|
for data in extracted_data:
|
|
que.put(data)
|
|
print(f"Total samples added to request queue: {len(extracted_data)}")
|
|
return que
|
|
|
|
|
|
def sample_requests_inner(dataset_path: str, num_requests: int):
|
|
start_time = time.perf_counter()
|
|
dataset = []
|
|
prompts_token_lens = []
|
|
completions_token_lens = []
|
|
text_test = []
|
|
with open(dataset_path) as f:
|
|
for line in f:
|
|
data = json.loads(line)
|
|
prompts_token_lens.append(int(data["input_token_num"]))
|
|
completions_token_lens.append(int(data["min_dec_len"]) - 1)
|
|
dataset.append([data["input_ids"], ""])
|
|
text_test.append(data["text"])
|
|
|
|
print(
|
|
f"Input length range: [{min(prompts_token_lens)}, {max(prompts_token_lens)}], "
|
|
f"Output length range: [{min(completions_token_lens)}, {max(completions_token_lens)}]"
|
|
)
|
|
avg_input_len = sum(prompts_token_lens) / len(prompts_token_lens)
|
|
avg_output_len = sum(completions_token_lens) / len(completions_token_lens)
|
|
|
|
print(f"Average input length: {avg_input_len:.2f}")
|
|
print(f"Average output length: {avg_output_len:.2f}")
|
|
print(f"Elapsed time for tokenization: {round(time.perf_counter() - start_time, 2)} seconds")
|
|
|
|
start_time = time.perf_counter()
|
|
filtered_dataset = []
|
|
for (prompt, _), input_len, output_len, text in zip(
|
|
dataset, prompts_token_lens, completions_token_lens, text_test
|
|
):
|
|
filtered_dataset.append([text, input_len, output_len])
|
|
|
|
sampled_requests = random.sample(filtered_dataset, num_requests)
|
|
print(f"Number of sampled requests: {len(sampled_requests)}")
|
|
que = mp.Queue()
|
|
for data in sampled_requests:
|
|
que.put(data)
|
|
print(f"Elapsed time for filtering: {round(time.perf_counter() - start_time, 2)} seconds")
|
|
return que
|
|
|
|
def save_results_to_file(stats: np.ndarray, output_file: str):
|
|
"""Save the benchmark results to a file."""
|
|
results = {
|
|
"first_token_latency": {
|
|
"min": stats[:, 0].min(),
|
|
"max": stats[:, 0].max(),
|
|
"avg": stats[:, 0].mean(),
|
|
},
|
|
"total_token_latency": {
|
|
"min": stats[:, 1].min(),
|
|
"max": stats[:, 1].max(),
|
|
"avg": stats[:, 1].mean(),
|
|
},
|
|
"input_sequence_length": {
|
|
"min": stats[:, 2].min(),
|
|
"max": stats[:, 2].max(),
|
|
"avg": stats[:, 2].mean(),
|
|
},
|
|
"output_sequence_length": {
|
|
"min": stats[:, 3].min(),
|
|
"max": stats[:, 3].max(),
|
|
"avg": stats[:, 3].mean(),
|
|
},
|
|
"real_output_sequence_length": {
|
|
"min": stats[:, 4].min(),
|
|
"max": stats[:, 4].max(),
|
|
"avg": stats[:, 4].mean(),
|
|
},
|
|
"inference_first_token_latency": {
|
|
"min": stats[:, 5].min(),
|
|
"max": stats[:, 5].max(),
|
|
"avg": stats[:, 5].mean(),
|
|
},
|
|
"inference_total_token_latency": {
|
|
"min": stats[:, 6].min(),
|
|
"max": stats[:, 6].max(),
|
|
"avg": stats[:, 6].mean(),
|
|
},
|
|
"qps": len(stats) / (stats[:, 1].sum() / len(stats)),
|
|
"real_output_tokens_per_second": len(stats) / (stats[:, 1].sum() / len(stats)) * stats[:, 4].mean(),
|
|
}
|
|
|
|
with open(output_file, "w") as f:
|
|
json.dump(results, f, indent=4)
|
|
print(f"Results saved to {output_file}")
|
|
|
|
|
|
def main(
|
|
dataset_path: str = "./filtered_sharedgpt_short_3000.json",
|
|
concurrency: int = 256,
|
|
num_prompts: int = 3000,
|
|
backend: str = "paddle",
|
|
model_name: str = "openlm-research/open_llama_13b",
|
|
host: str = "localhost",
|
|
port: str = "8100",
|
|
warmup_round: int = 1,
|
|
dataset_name: str = "sharegpt",
|
|
min_dec_len: int = 1,
|
|
max_dec_len: int = 2048,
|
|
output_file: Optional[str] = None,
|
|
):
|
|
api_url = f"http://{host}:{port}/generate_stream"
|
|
if backend in ["vllm", "trtllm", "paddle"]:
|
|
api_url = f"http://{host}:{port}/v1/chat/completions"
|
|
elif backend == "sglang":
|
|
api_url = f"http://{host}:{port}/generate"
|
|
else:
|
|
raise RuntimeError("Unsupported backend. Choose from: paddle, vllm, trtllm, sglang")
|
|
|
|
print(f"API URL: {api_url}")
|
|
warmup(concurrency, 256, warmup_round, backend, api_url, model_name)
|
|
print(f"Dataset name: {dataset_name}")
|
|
|
|
if dataset_name == "sharegpt":
|
|
req_que = sample_requests_filtered_shared_gpt(dataset_path, num_prompts)
|
|
elif dataset_name == "paddle_inner":
|
|
req_que = sample_requests_inner(dataset_path, num_prompts)
|
|
else:
|
|
raise ValueError("Invalid dataset name. Choose from: sharegpt, paddle_inner")
|
|
|
|
res_que = mp.Queue()
|
|
procs = []
|
|
end_flags = mp.Queue()
|
|
start_time = time.perf_counter()
|
|
for i in range(concurrency):
|
|
proc = mp.Process(
|
|
target=infer,
|
|
args=(i + 1, req_que, res_que, end_flags, backend, api_url, model_name, min_dec_len, max_dec_len),
|
|
)
|
|
procs.append(proc)
|
|
proc.start()
|
|
|
|
while end_flags.qsize() < concurrency:
|
|
time.sleep(0.01)
|
|
|
|
end_time = time.perf_counter()
|
|
elapsed_time = end_time - start_time
|
|
|
|
stats = [[]] * (concurrency + 1)
|
|
while res_que.qsize() > 0:
|
|
session_id, _stats = res_que.get()
|
|
stats[session_id] = np.array(_stats)
|
|
|
|
all_stat = []
|
|
for i in range(concurrency):
|
|
print(f"Session {i + 1} processed {len(stats[i + 1])} requests.")
|
|
if len(stats[i + 1]) > 0:
|
|
all_stat.append(stats[i + 1])
|
|
|
|
all_stat = np.concatenate(all_stat).reshape(-1, 7)
|
|
|
|
print(f"Processed {len(all_stat)} requests in {elapsed_time:.2f} seconds.")
|
|
print(
|
|
f"First Token Latency (min, max, avg): {all_stat[:, 0].min():.2f}, {all_stat[:, 0].max():.2f}, {all_stat[:, 0].mean():.2f}"
|
|
)
|
|
print(
|
|
f"Total Token Latency (min, max, avg): {all_stat[:, 1].min():.2f}, {all_stat[:, 1].max():.2f}, {all_stat[:, 1].mean():.2f}"
|
|
)
|
|
print(
|
|
f"Input Sequence Length (min, max, avg): {all_stat[:, 2].min():.2f}, {all_stat[:, 2].max():.2f}, {all_stat[:, 2].mean():.2f}"
|
|
)
|
|
print(
|
|
f"Output Sequence Length (min, max, avg): {all_stat[:, 3].min():.2f}, {all_stat[:, 3].max():.2f}, {all_stat[:, 3].mean():.2f}"
|
|
)
|
|
print(
|
|
f"Real Output Sequence Length (min, max, avg): {all_stat[:, 4].min():.2f}, {all_stat[:, 4].max():.2f}, {all_stat[:, 4].mean():.2f}"
|
|
)
|
|
print(
|
|
f"Inference First Token Latency (min, max, avg): {all_stat[:, 5].min():.2f}, {all_stat[:, 5].max():.2f}, {all_stat[:, 5].mean():.2f}"
|
|
)
|
|
print(
|
|
f"Inference Total Token Latency (min, max, avg): {all_stat[:, 6].min():.2f}, {all_stat[:, 6].max():.2f}, {all_stat[:, 6].mean():.2f}"
|
|
)
|
|
print(f"QPS: {len(all_stat) / elapsed_time:.2f}")
|
|
print(f"Real Output Tokens/s: {len(all_stat) / elapsed_time * all_stat[:, 4].mean():.2f}")
|
|
|
|
# Save results to file if output_file is provided
|
|
if output_file:
|
|
save_results_to_file(all_stat, output_file)
|
|
|
|
for p in procs:
|
|
p.join()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Benchmark script for inference performance.")
|
|
parser.add_argument("--dataset_path", type=str, default="./filtered_sharedgpt_short_3000.json", help="Path to the dataset file.")
|
|
parser.add_argument("--concurrency", type=int, default=256, help="Number of concurrent requests.")
|
|
parser.add_argument("--num_prompts", type=int, default=3000, help="Number of prompts to process.")
|
|
parser.add_argument("--backend", type=str, default="paddle", choices=["paddle", "vllm", "trtllm", "sglang"], help="Backend to use for inference.")
|
|
parser.add_argument("--model_name", type=str, default="openlm-research/open_llama_13b", help="Name of the model to use.")
|
|
parser.add_argument("--host", type=str, default="localhost", help="Host address of the inference server.")
|
|
parser.add_argument("--port", type=str, default="8100", help="Port of the inference server.")
|
|
parser.add_argument("--warmup_round", type=int, default=1, help="Number of warmup rounds.")
|
|
parser.add_argument("--dataset_name", type=str, default="sharegpt", choices=["sharegpt", "paddle_inner"], help="Name of the dataset to use.")
|
|
parser.add_argument("--min_dec_len", type=int, default=1, help="Minimum decoding length.")
|
|
parser.add_argument("--max_dec_len", type=int, default=2048, help="Maximum decoding length.")
|
|
parser.add_argument("--output_file", type=str, default=None, help="Path to save the results file.")
|
|
|
|
args = parser.parse_args()
|
|
main(
|
|
dataset_path=args.dataset_path,
|
|
concurrency=args.concurrency,
|
|
num_prompts=args.num_prompts,
|
|
backend=args.backend,
|
|
model_name=args.model_name,
|
|
host=args.host,
|
|
port=args.port,
|
|
warmup_round=args.warmup_round,
|
|
dataset_name=args.dataset_name,
|
|
min_dec_len=args.min_dec_len,
|
|
max_dec_len=args.max_dec_len,
|
|
output_file=args.output_file,
|
|
) |