539 lines
21 KiB
Python
539 lines
21 KiB
Python
import argparse
|
|
import asyncio
|
|
import json
|
|
import queue
|
|
import random
|
|
import threading
|
|
import time
|
|
from datetime import datetime
|
|
|
|
import numpy as np
|
|
import requests
|
|
from tqdm.asyncio import tqdm
|
|
|
|
from sglang.benchmark.datasets import sample_random_requests
|
|
from sglang.benchmark.utils import get_tokenizer
|
|
from sglang.test.kits.cache_hit_kit import async_request_sglang_generate, gen_payload
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description="Script to benchmark concurrent requests to a server."
|
|
)
|
|
parser.add_argument(
|
|
"--num-clients",
|
|
type=int,
|
|
default=256,
|
|
help="Number of concurrent clients",
|
|
)
|
|
parser.add_argument(
|
|
"--max-parallel",
|
|
type=int,
|
|
default=128,
|
|
help="Maximum number of parallel requests",
|
|
)
|
|
parser.add_argument(
|
|
"--request-length",
|
|
type=int,
|
|
default=512,
|
|
help="Length of each new request",
|
|
)
|
|
parser.add_argument(
|
|
"--output-length",
|
|
type=int,
|
|
default=64,
|
|
help="Length of each output",
|
|
)
|
|
parser.add_argument(
|
|
"--num-rounds",
|
|
type=int,
|
|
default=5,
|
|
help="Number of rounds per client",
|
|
)
|
|
parser.add_argument(
|
|
"--distribution",
|
|
type=str,
|
|
default="poisson",
|
|
choices=["poisson", "uniform"],
|
|
help="Distribution type for request intervals (poisson or uniform)",
|
|
)
|
|
parser.add_argument(
|
|
"--request-rate",
|
|
type=float,
|
|
default=1.0,
|
|
help="Average number of requests per second",
|
|
)
|
|
parser.add_argument(
|
|
"--host",
|
|
type=str,
|
|
default="localhost",
|
|
help="Server hostname or IP (default: localhost)",
|
|
)
|
|
parser.add_argument(
|
|
"--port",
|
|
type=int,
|
|
default=30000,
|
|
help="Server port (default: 30000)",
|
|
)
|
|
parser.add_argument(
|
|
"--model-path",
|
|
type=str,
|
|
default="meta-llama/Llama-3.1-8B-Instruct",
|
|
help="model path compatible with Hugging Face Transformers",
|
|
)
|
|
parser.add_argument(
|
|
"--dataset-path",
|
|
type=str,
|
|
default="",
|
|
help="local dataset to sample tokens from",
|
|
)
|
|
parser.add_argument(
|
|
"--log-file",
|
|
type=str,
|
|
default="performance_metrics.jsonl",
|
|
help="File to log performance metrics",
|
|
)
|
|
parser.add_argument(
|
|
"--disable-auto-run",
|
|
action="store_true",
|
|
help="If set, disable automatically testing with a range of request rates.",
|
|
)
|
|
parser.add_argument(
|
|
"--disable-random-sample",
|
|
action="store_true",
|
|
help="If set, disable random sampling of requests from the ShareGPT dataset.",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-round-barrier",
|
|
action="store_true",
|
|
help="If set, only send i-th turn requests after all (i-1)-th turn requests finished.",
|
|
)
|
|
parser.add_argument(
|
|
"--sub-question-input-length",
|
|
type=int,
|
|
default=0,
|
|
help="Length of the sub question input for each request, if set 0 use request_length",
|
|
)
|
|
parser.add_argument(
|
|
"--ready-queue-policy",
|
|
type=str,
|
|
default="random",
|
|
help="Policy for popping requests from the ready queue (random or fifo)",
|
|
)
|
|
parser.add_argument(
|
|
"--tag",
|
|
type=str,
|
|
default="",
|
|
help="Tag of a certain run in the log file",
|
|
)
|
|
parser.add_argument("--seed", type=int, default=1, help="The random seed.")
|
|
parser.add_argument(
|
|
"--lora-path",
|
|
type=str,
|
|
default="",
|
|
help="String of LoRA path. Currently we only support benchmarking on a single LoRA adaptor.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def log_to_jsonl_file(data, file_path="performance_metrics.jsonl", tag=""):
|
|
"""Append the data with a timestamp and tag to the specified JSONL file."""
|
|
timestamped_data = {"timestamp": datetime.now().isoformat(), "tag": tag, **data}
|
|
try:
|
|
with open(file_path, "a") as file:
|
|
file.write(
|
|
json.dumps(timestamped_data) + "\n"
|
|
) # Write as a single line in JSONL format
|
|
except IOError as e:
|
|
print(f"Error writing to JSONL file: {e}")
|
|
|
|
|
|
class ReadyQueue:
|
|
"""
|
|
Thread-safe queue that can pop requests in different orders based on given policy.
|
|
"""
|
|
|
|
def __init__(self, init_requests=None, policy="random"):
|
|
self.lock = threading.Lock()
|
|
self.requests = init_requests or []
|
|
self.policy = policy
|
|
|
|
def append(self, item):
|
|
with self.lock:
|
|
self.requests.append(item)
|
|
|
|
def pop(self):
|
|
with self.lock:
|
|
if not self.requests:
|
|
return None
|
|
if self.policy == "random":
|
|
index = random.randrange(len(self.requests))
|
|
return self.requests.pop(index)
|
|
elif self.policy == "fifo":
|
|
return self.requests.pop(0)
|
|
else:
|
|
# todo, varying thinking time of clients
|
|
raise ValueError(f"{self.policy} not implemented")
|
|
|
|
|
|
class WorkloadGenerator:
|
|
def __init__(self, args):
|
|
# Construct the base URL for requests
|
|
self.url = f"http://{args.host}:{args.port}/generate"
|
|
|
|
self.tokenizer = get_tokenizer(args.model_path)
|
|
self.distribution = args.distribution
|
|
self.request_rate = args.request_rate
|
|
self.start_time = None
|
|
self.finished_time = None
|
|
|
|
self.sent_requests = 0
|
|
self.completed_requests = 0
|
|
|
|
# Use return_text=False to get token ids instead of text
|
|
self.candidate_inputs = sample_random_requests(
|
|
input_len=args.request_length,
|
|
output_len=args.output_length,
|
|
num_prompts=args.num_clients,
|
|
range_ratio=1.0,
|
|
tokenizer=self.tokenizer,
|
|
dataset_path=args.dataset_path,
|
|
random_sample=not args.disable_random_sample,
|
|
return_text=False,
|
|
)
|
|
# r.prompt is now List[int] when return_text=False
|
|
self.candidate_inputs = [list(i.prompt) for i in self.candidate_inputs]
|
|
|
|
if args.sub_question_input_length != 0:
|
|
sub_question_input_length = args.sub_question_input_length
|
|
else:
|
|
sub_question_input_length = args.request_length
|
|
|
|
self.sub_question_inputs = sample_random_requests(
|
|
input_len=sub_question_input_length,
|
|
output_len=args.output_length,
|
|
num_prompts=args.num_clients * max(args.num_rounds - 1, 1),
|
|
range_ratio=1.0,
|
|
tokenizer=self.tokenizer,
|
|
dataset_path=args.dataset_path,
|
|
random_sample=not args.disable_random_sample,
|
|
return_text=False,
|
|
)
|
|
|
|
init_requests = [
|
|
(
|
|
i,
|
|
gen_payload(
|
|
self.candidate_inputs[i], args.output_length, args.lora_path
|
|
),
|
|
)
|
|
for i in range(args.num_clients)
|
|
]
|
|
# history now stores List[int] (token ids) for each client
|
|
self.client_records = {
|
|
i: {"round": 0, "history": list(self.candidate_inputs[i])}
|
|
for i in range(args.num_clients)
|
|
}
|
|
self.ready_queue = ReadyQueue(
|
|
init_requests=init_requests, policy=args.ready_queue_policy
|
|
)
|
|
self.candidate_inputs = self.candidate_inputs[args.num_clients :]
|
|
|
|
self.response_queue = queue.Queue()
|
|
self.pbar = tqdm(total=args.num_clients * args.num_rounds)
|
|
self.performance_metrics = {
|
|
"ttft": [],
|
|
"latency": [],
|
|
"prompt_len": [],
|
|
"cached_tokens": [],
|
|
"generated_len": [],
|
|
}
|
|
self.enable_round_barrier = args.enable_round_barrier
|
|
if self.enable_round_barrier:
|
|
# Add round-specific metrics while preserving the original structure
|
|
for i in range(args.num_rounds):
|
|
self.performance_metrics[f"round_{i}"] = {
|
|
"ttft": [],
|
|
"latency": [],
|
|
"prompt_len": [],
|
|
"cached_tokens": [],
|
|
"generated_len": [],
|
|
}
|
|
self.num_clients = args.num_clients
|
|
|
|
self.num_rounds = args.num_rounds
|
|
self.max_parallel = args.max_parallel
|
|
self.output_length = args.output_length
|
|
|
|
async def handle_request(self, item):
|
|
try:
|
|
client_id, payload = item
|
|
response = await async_request_sglang_generate(payload, self.url, self.pbar)
|
|
if self.pbar.n == self.pbar.total:
|
|
self.finished_time = time.perf_counter()
|
|
self.response_queue.put((client_id, response))
|
|
except Exception as e:
|
|
print(f"Request failed: {e}")
|
|
|
|
def request_sender(self):
|
|
async def request_loop():
|
|
while True:
|
|
if self.sent_requests - self.completed_requests < self.max_parallel:
|
|
new_request = self.ready_queue.pop()
|
|
if new_request:
|
|
asyncio.create_task(self.handle_request(new_request))
|
|
self.sent_requests += 1
|
|
else:
|
|
await asyncio.sleep(0.05)
|
|
continue
|
|
|
|
if self.pbar.n == self.pbar.total:
|
|
break
|
|
|
|
# Calculate Poisson-distributed wait time
|
|
if self.distribution == "poisson":
|
|
sleep_time = random.expovariate(self.request_rate)
|
|
elif self.distribution == "uniform":
|
|
avg_interval = (
|
|
1.0 / self.request_rate if self.request_rate > 0 else 1.0
|
|
)
|
|
sleep_time = random.uniform(0, 2 * avg_interval)
|
|
else:
|
|
raise ValueError("Invalid distribution type")
|
|
await asyncio.sleep(sleep_time) # Wait before sending the next request
|
|
|
|
# Create and run the event loop for asynchronous requests
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
loop.run_until_complete(request_loop())
|
|
loop.close()
|
|
|
|
def response_handler(self):
|
|
next_round_reqs = []
|
|
while True:
|
|
try:
|
|
client_id, response = self.response_queue.get(
|
|
timeout=10
|
|
) # Block until response is available
|
|
if not response.success:
|
|
raise ValueError(f"Request failed with error: {response.error}")
|
|
# Use output_ids (token ids) instead of generated_text
|
|
self.client_records[client_id]["history"].extend(response.output_ids)
|
|
current_round = self.client_records[client_id]["round"]
|
|
self.client_records[client_id]["round"] += 1
|
|
self.performance_metrics["ttft"].append(response.ttft)
|
|
self.performance_metrics["latency"].append(response.latency)
|
|
self.performance_metrics["prompt_len"].append(response.prompt_len)
|
|
self.performance_metrics["cached_tokens"].append(response.cached_tokens)
|
|
self.performance_metrics["generated_len"].append(response.generated_len)
|
|
if self.enable_round_barrier:
|
|
self.performance_metrics[f"round_{current_round}"]["ttft"].append(
|
|
response.ttft
|
|
)
|
|
self.performance_metrics[f"round_{current_round}"][
|
|
"latency"
|
|
].append(response.latency)
|
|
self.performance_metrics[f"round_{current_round}"][
|
|
"prompt_len"
|
|
].append(response.prompt_len)
|
|
self.performance_metrics[f"round_{current_round}"][
|
|
"cached_tokens"
|
|
].append(response.cached_tokens)
|
|
self.performance_metrics[f"round_{current_round}"][
|
|
"generated_len"
|
|
].append(response.generated_len)
|
|
self.completed_requests += 1
|
|
|
|
if self.client_records[client_id]["round"] < self.num_rounds:
|
|
# Append sub-question token ids to client's history
|
|
sub_q_ids = list(self.sub_question_inputs.pop().prompt)
|
|
self.client_records[client_id]["history"].extend(sub_q_ids)
|
|
new_req = (
|
|
client_id,
|
|
gen_payload(
|
|
self.client_records[client_id]["history"],
|
|
self.output_length,
|
|
args.lora_path,
|
|
),
|
|
)
|
|
if self.enable_round_barrier:
|
|
next_round_reqs.append(new_req)
|
|
if len(next_round_reqs) == self.num_clients:
|
|
for req in next_round_reqs:
|
|
self.ready_queue.append(req)
|
|
next_round_reqs = []
|
|
else:
|
|
self.ready_queue.append(new_req)
|
|
except queue.Empty:
|
|
if self.pbar.n == self.pbar.total:
|
|
break
|
|
except ValueError as e:
|
|
print(f"Error processing response for client {client_id}: {e}")
|
|
continue
|
|
|
|
def run(self):
|
|
request_thread = threading.Thread(target=self.request_sender, daemon=True)
|
|
response_thread = threading.Thread(target=self.response_handler, daemon=True)
|
|
|
|
self.start_time = time.perf_counter()
|
|
request_thread.start()
|
|
response_thread.start()
|
|
|
|
request_thread.join()
|
|
response_thread.join()
|
|
self.pbar.close()
|
|
|
|
duration = self.finished_time - self.start_time
|
|
sorted_ttft = sorted(self.performance_metrics["ttft"])
|
|
sorted_latency = sorted(self.performance_metrics["latency"])
|
|
|
|
def percentile(sorted_vals, q):
|
|
if not sorted_vals:
|
|
return 0.0
|
|
idx = int(q * len(sorted_vals))
|
|
if idx >= len(sorted_vals):
|
|
idx = len(sorted_vals) - 1
|
|
return sorted_vals[idx]
|
|
|
|
def max_or_zero(sorted_vals):
|
|
return sorted_vals[-1] if sorted_vals else 0.0
|
|
|
|
performance_data = {
|
|
"summary": {
|
|
"total_requests": len(self.performance_metrics["ttft"]),
|
|
"request_rate": self.request_rate,
|
|
"average_prompt_len": (
|
|
sum(self.performance_metrics["prompt_len"])
|
|
/ len(self.performance_metrics["prompt_len"])
|
|
if self.performance_metrics["prompt_len"]
|
|
else 0.0
|
|
),
|
|
"average_output_len": (
|
|
sum(self.performance_metrics["generated_len"])
|
|
/ len(self.performance_metrics["generated_len"])
|
|
if self.performance_metrics["generated_len"]
|
|
else 0.0
|
|
),
|
|
"average_ttft": sum(self.performance_metrics["ttft"])
|
|
/ len(self.performance_metrics["ttft"]),
|
|
"p90_ttft": percentile(sorted_ttft, 0.9),
|
|
"p99_ttft": percentile(sorted_ttft, 0.99),
|
|
"median_ttft": percentile(sorted_ttft, 0.5),
|
|
"max_ttft": max_or_zero(sorted_ttft),
|
|
"average_latency": sum(self.performance_metrics["latency"])
|
|
/ len(self.performance_metrics["latency"]),
|
|
"p90_latency": percentile(sorted_latency, 0.9),
|
|
"p99_latency": percentile(sorted_latency, 0.99),
|
|
"median_latency": percentile(sorted_latency, 0.5),
|
|
"max_latency": max_or_zero(sorted_latency),
|
|
"input_token_throughput": sum(self.performance_metrics["prompt_len"])
|
|
/ duration,
|
|
"output_token_throughput": sum(
|
|
self.performance_metrics["generated_len"]
|
|
)
|
|
/ duration,
|
|
"throughput": self.pbar.total / duration,
|
|
"cache_hit_rate": (
|
|
0
|
|
if sum(self.performance_metrics["prompt_len"]) == 0
|
|
else sum(self.performance_metrics["cached_tokens"])
|
|
/ sum(self.performance_metrics["prompt_len"])
|
|
),
|
|
},
|
|
}
|
|
if self.enable_round_barrier:
|
|
performance_data["round"] = {}
|
|
for round_num in range(args.num_rounds):
|
|
round_key = f"round_{round_num}"
|
|
round_metrics = self.performance_metrics[round_key]
|
|
performance_data["round"][round_key] = {
|
|
"average_ttft": (
|
|
sum(round_metrics["ttft"]) / len(round_metrics["ttft"])
|
|
if round_metrics["ttft"]
|
|
else 0
|
|
),
|
|
"cache_hit_rate": (
|
|
0
|
|
if sum(round_metrics["prompt_len"]) == 0
|
|
else sum(round_metrics["cached_tokens"])
|
|
/ sum(round_metrics["prompt_len"])
|
|
),
|
|
"request_count": len(round_metrics["ttft"]),
|
|
}
|
|
print("All requests completed")
|
|
print("Performance metrics summary:")
|
|
print(
|
|
f" Total requests: {performance_data['summary']['total_requests']} at {performance_data['summary']['request_rate']} requests per second"
|
|
)
|
|
print(
|
|
f" Average Prompt Length: {performance_data['summary']['average_prompt_len']:.2f} tokens"
|
|
)
|
|
print(
|
|
f" Average Output Length: {performance_data['summary']['average_output_len']:.2f} tokens"
|
|
)
|
|
print(f" Average TTFT: {performance_data['summary']['average_ttft']:.2f}")
|
|
print(f" P90 TTFT: {performance_data['summary']['p90_ttft']:.2f}")
|
|
print(f" P99 TTFT: {performance_data['summary']['p99_ttft']:.2f}")
|
|
print(f" Median TTFT: {performance_data['summary']['median_ttft']:.2f}")
|
|
print(f" Max TTFT: {performance_data['summary']['max_ttft']:.2f}")
|
|
print(
|
|
f" Average latency: {performance_data['summary']['average_latency']:.2f}"
|
|
)
|
|
print(f" P90 latency: {performance_data['summary']['p90_latency']:.2f}")
|
|
print(f" P99 latency: {performance_data['summary']['p99_latency']:.2f}")
|
|
print(f" Median latency: {performance_data['summary']['median_latency']:.2f}")
|
|
print(f" Max latency: {performance_data['summary']['max_latency']:.2f}")
|
|
print(
|
|
f" Input token throughput: {performance_data['summary']['input_token_throughput']:.2f} tokens per second"
|
|
)
|
|
print(
|
|
f" Output token throughput: {performance_data['summary']['output_token_throughput']:.2f} tokens per second"
|
|
)
|
|
print(
|
|
f" Request Throughput: {performance_data['summary']['throughput']:.2f} requests per second"
|
|
)
|
|
print(f" Cache Hit Rate: {performance_data['summary']['cache_hit_rate']:.6f}")
|
|
|
|
if self.enable_round_barrier:
|
|
# Print round-basedsummary
|
|
print("Per-round metrics:")
|
|
if "round" in performance_data:
|
|
for round_num in range(self.num_rounds):
|
|
round_key = f"round_{round_num}"
|
|
if round_key in performance_data["round"]:
|
|
round_data = performance_data["round"][round_key]
|
|
avg_ttft = round_data["average_ttft"]
|
|
cache_hit_rate = round_data["cache_hit_rate"]
|
|
request_count = round_data["request_count"]
|
|
print(
|
|
f" Round {round_num}: Average TTFT = {avg_ttft:.2f}s, "
|
|
f"Cache Hit Rate = {cache_hit_rate:.6f} "
|
|
f"({request_count} requests)"
|
|
)
|
|
else:
|
|
print(f" Round {round_num}: No requests completed")
|
|
|
|
return performance_data
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
flush_cache_url = f"http://{args.host}:{args.port}/flush_cache"
|
|
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
|
|
if args.disable_auto_run:
|
|
print("Running with specified request rate...")
|
|
request_rates = [args.request_rate]
|
|
else:
|
|
print("Auto-running with different request rates...")
|
|
request_rates = [16, 14, 12, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
|
|
|
|
for rate in request_rates:
|
|
args.request_rate = rate
|
|
requests.post(flush_cache_url)
|
|
time.sleep(1)
|
|
performance_data = WorkloadGenerator(args).run()
|
|
log_to_jsonl_file(performance_data, args.log_file, tag=args.tag)
|