wehub-resource-sync
779 lines
24 KiB
Python
779 lines
24 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
# Adapted from
|
||
# https://github.com/vllm-project/vllm/blob/main/benchmarks/benchmark_long_document_qa_throughput.py
|
||
|
||
"""
|
||
Commandline arguments:
|
||
--num-documents: The number of documents to sample prompts from.
|
||
|
||
--document-length: The length of each document in tokens.
|
||
(Optional, default: 20000)
|
||
|
||
--output-len: The number of tokens to generate for each prompt.
|
||
(Optional, default: 100)
|
||
|
||
--repeat-count: The number of times to repeat each prompt.
|
||
(Optional, default: 2)
|
||
|
||
--repeat-mode: The mode to repeat prompts. The supported modes are:
|
||
- 'random': shuffle the prompts randomly. (Default)
|
||
- 'tile': the entire prompt list is repeated in sequence.
|
||
- 'interleave': each prompt is repeated consecutively before
|
||
moving to the next element.
|
||
|
||
--shuffle-seed: Random seed when the repeat mode is "random".
|
||
(Optional, default: 0)
|
||
|
||
--host: Host to query the vLLM server
|
||
--port: Port to query the vLLM server
|
||
--base-url: Base URL to query the LLM server (exclusive with --host/--port)
|
||
|
||
--model: Model name
|
||
|
||
--max-inflight-requests: Maximum number of in-flight requests. Default is 2
|
||
|
||
--sleep-time-after-warmup: Sleep time after warm up iteration.
|
||
(Optional, default: 0.0 seconds)
|
||
|
||
--output: Filename to write all responses to. If omitted, writes to stdout.
|
||
|
||
--expected-ttft-gain: Expected minimum speed-up in time-to-first-token
|
||
(warmup/query) as a factor, e.g. 4.3 for 4.3×. If
|
||
actual gain is below this, exits.
|
||
|
||
--expected-latency-gain: Expected minimum speed-up in total round time
|
||
(warmup/query) as a factor, e.g. 4.5 for 4.5×.
|
||
If actual gain is below this, exits.
|
||
|
||
--expected-latency: Expected end to end latency for the first query round.
|
||
--completions: Use completions API instead of chat completions API
|
||
|
||
--visualize: Visualize the results
|
||
|
||
--eos-token-id: EOS token id. we bias against this token id so we always
|
||
get the number of output tokens we specify
|
||
|
||
--hit-miss-ratio: In query round, control how many of the prompts
|
||
will miss the cache. For example, 3:1 means every fourth repeated prompt
|
||
will miss the cache. 2:2 means 2 hits and 2 misses.
|
||
|
||
--trim-fraction: Exclude the smallest X fraction and largest X fraction
|
||
and calculate mean of the rest. For example, 0.1 drops bottom 10% and top 10%.
|
||
"""
|
||
|
||
# Standard
|
||
from dataclasses import dataclass
|
||
import argparse
|
||
import asyncio
|
||
import os
|
||
import random
|
||
import sys
|
||
import time
|
||
|
||
# Third Party
|
||
from openai import AsyncOpenAI
|
||
import pandas as pd
|
||
|
||
# Global output filename (set in __main__)
|
||
OUTPUT_FILE = None
|
||
completions_mode = False
|
||
visualize = False
|
||
eos_token_id = None
|
||
|
||
|
||
@dataclass
|
||
class RequestStats:
|
||
prompt_id: int
|
||
request_start: float
|
||
ttft: float
|
||
request_end: float
|
||
successful: bool
|
||
|
||
|
||
def get_url_from_args(args):
|
||
"""
|
||
Get the base URL from command line arguments.
|
||
Args:
|
||
args: Command line arguments.
|
||
Returns:
|
||
str: The base URL.
|
||
"""
|
||
if args.base_url is not None:
|
||
return args.base_url
|
||
else:
|
||
host = args.host if args.host is not None else "localhost"
|
||
port = args.port if args.port is not None else 8000
|
||
return f"http://{host}:{port}/v1"
|
||
|
||
|
||
def extract_reasoning_content(chunk):
|
||
"""
|
||
Extract reasoning content from the response chunk.
|
||
Args:
|
||
chunk: The response chunk from OpenAI Chat Completions API.
|
||
Returns:
|
||
str | None: The reasoning content extracted from the chunk.
|
||
None means no reasoning content in this chunk.
|
||
"""
|
||
delta = chunk.choices[0].delta
|
||
potential_reasoning_keys = [
|
||
"reasoning_content",
|
||
"reasoning",
|
||
"tool_calls",
|
||
"tool_call",
|
||
"tool_responses",
|
||
]
|
||
for key in potential_reasoning_keys:
|
||
if hasattr(delta, key) and getattr(delta, key):
|
||
return getattr(delta, key)
|
||
return None
|
||
|
||
|
||
def extract_normal_content(chunk):
|
||
"""
|
||
Extract normal content from the response chunk.
|
||
Args:
|
||
chunk: The response chunk from OpenAI Chat Completions API.
|
||
Returns:
|
||
str | None: The normal content extracted from the chunk.
|
||
None means no normal content in this chunk.
|
||
"""
|
||
delta = chunk.choices[0].delta
|
||
if hasattr(delta, "content") and delta.content:
|
||
return chunk.choices[0].delta.content
|
||
return None
|
||
|
||
|
||
def has_content_completions(chunk):
|
||
"""
|
||
Completions streaming emits text at choices[0].text.
|
||
"""
|
||
return bool(chunk.choices) and (chunk.choices[0].text is not None)
|
||
|
||
|
||
def has_content(chunk, completions_mode=False):
|
||
"""
|
||
Check if the chunk has content in the choices.
|
||
Args:
|
||
chunk: The response chunk from OpenAI Chat Completions API.
|
||
|
||
Returns:
|
||
bool: True if content exists, False otherwise.
|
||
"""
|
||
if completions_mode:
|
||
return has_content_completions(chunk)
|
||
|
||
return (
|
||
chunk.choices
|
||
and chunk.choices[0].delta
|
||
and (
|
||
extract_normal_content(chunk) is not None
|
||
or extract_reasoning_content(chunk) is not None
|
||
)
|
||
)
|
||
|
||
|
||
def extract_content_completions(chunk):
|
||
"""
|
||
Extract content from a Completions stream chunk.
|
||
"""
|
||
return chunk.choices[0].text or ""
|
||
|
||
|
||
def extract_content(chunk, completions_mode=False):
|
||
"""
|
||
Extract content from the response chunk.
|
||
Args:
|
||
chunk: The response chunk from OpenAI Chat Completions API.
|
||
Returns:
|
||
str: The content extracted from the chunk.
|
||
"""
|
||
if completions_mode:
|
||
return extract_content_completions(chunk)
|
||
|
||
if normal_content := extract_normal_content(chunk):
|
||
return normal_content
|
||
elif reasoning_content := extract_reasoning_content(chunk):
|
||
return reasoning_content
|
||
else:
|
||
return ""
|
||
|
||
|
||
def write_resp(text: str):
|
||
"""
|
||
Write text to the specified output file (if any), otherwise to stdout.
|
||
"""
|
||
if OUTPUT_FILE:
|
||
with open(OUTPUT_FILE, "a") as resp_file:
|
||
resp_file.write(text)
|
||
else:
|
||
sys.stdout.write(text)
|
||
|
||
|
||
async def process_single_prompt(
|
||
client, model, prompt, prompt_index, total_prompts, output_len, semaphore
|
||
) -> RequestStats:
|
||
"""
|
||
Process a single prompt with the given client and model.
|
||
|
||
Args:
|
||
client: The OpenAI client for making API calls.
|
||
model: The model name to use for generation.
|
||
prompt: The prompt string to be processed.
|
||
prompt_index: Index of the current prompt (0-based).
|
||
total_prompts: Total number of prompts being processed.
|
||
output_len: The maximum number of tokens to generate.
|
||
semaphore: Asyncio semaphore to limit concurrent requests.
|
||
|
||
Returns:
|
||
RequestStats: RequestStats object containing the request stats
|
||
"""
|
||
async with semaphore: # Acquire semaphore to limit concurrent requests
|
||
write_resp(f"\n--- Sending prompt {prompt_index + 1}/{total_prompts} ---\n")
|
||
# a request starts once it acquires the semaphore
|
||
start_time = time.time()
|
||
first_token_time = None
|
||
|
||
# add stop None so we always get the number of output tokens we specify
|
||
if completions_mode:
|
||
response = await client.completions.create(
|
||
model=model,
|
||
prompt=prompt,
|
||
stream=True,
|
||
max_tokens=output_len,
|
||
temperature=0.0,
|
||
stream_options={"include_usage": True},
|
||
logit_bias={str(eos_token_id): -100}
|
||
if eos_token_id is not None
|
||
else None,
|
||
)
|
||
else:
|
||
response = await client.chat.completions.create(
|
||
model=model,
|
||
messages=[{"role": "user", "content": prompt}],
|
||
stream=True,
|
||
max_tokens=output_len,
|
||
temperature=0.0,
|
||
stream_options={"include_usage": True},
|
||
logit_bias={str(eos_token_id): -100}
|
||
if eos_token_id is not None
|
||
else None,
|
||
)
|
||
|
||
responses = []
|
||
# Collect the response chunks
|
||
async for chunk in response:
|
||
if not chunk.choices:
|
||
continue
|
||
|
||
# Handle content for chat completions
|
||
if has_content(chunk, completions_mode):
|
||
content = extract_content(chunk, completions_mode)
|
||
if first_token_time is None and content != "":
|
||
first_token_time = time.time()
|
||
responses.append(content)
|
||
|
||
end_time = time.time()
|
||
final_response = "".join(responses)
|
||
write_resp(f"\nResponse of request {prompt_index}: {final_response}\n")
|
||
|
||
# TTFT < 0 means not successful
|
||
ttft = (first_token_time - start_time) if first_token_time is not None else -1
|
||
return RequestStats(
|
||
prompt_id=prompt_index,
|
||
request_start=start_time,
|
||
ttft=ttft,
|
||
request_end=end_time,
|
||
successful=ttft > 0,
|
||
)
|
||
|
||
|
||
async def test_long_document_qa(
|
||
client, model, prompts=None, output_len=100, max_inflight_requests=10
|
||
) -> list[RequestStats]:
|
||
"""
|
||
Test long document QA with the given prompts and sampling parameters.
|
||
Process prompts concurrently with a limit on inflight requests.
|
||
|
||
Args:
|
||
client: The OpenAI client for making API calls.
|
||
model: The model name to use for generation.
|
||
prompts: A list of prompt strings to be processed by the LLM.
|
||
output_len: The maximum number of tokens to generate.
|
||
max_inflight_requests: Maximum number of concurrent requests.
|
||
|
||
Returns:
|
||
list: request_stats - a list of RequestStats objects
|
||
"""
|
||
# Create semaphore to limit concurrent requests
|
||
semaphore = asyncio.Semaphore(max_inflight_requests)
|
||
|
||
# Create tasks for all prompts
|
||
tasks = [
|
||
process_single_prompt(
|
||
client=client,
|
||
model=model,
|
||
prompt=prompt,
|
||
prompt_index=i,
|
||
total_prompts=len(prompts),
|
||
output_len=output_len,
|
||
semaphore=semaphore,
|
||
)
|
||
for i, prompt in enumerate(prompts)
|
||
]
|
||
# Execute all tasks concurrently and collect results
|
||
# The semaphore will control max concurrent requests
|
||
request_stats = await asyncio.gather(*tasks)
|
||
|
||
return request_stats
|
||
|
||
|
||
def repeat_prompts(prompts, repeat_count, mode: str):
|
||
"""
|
||
Repeat each prompt in the list for a specified number of times.
|
||
The order of prompts in the output list depends on the mode.
|
||
|
||
Args:
|
||
prompts: A list of prompts to be repeated.
|
||
repeat_count: The number of times each prompt is repeated.
|
||
mode: The mode of repetition. Supported modes are:
|
||
- 'random': Shuffle the prompts randomly after repetition.
|
||
- 'tile': Repeat the entire prompt list in sequence.
|
||
Example: [1, 2, 3] -> [1, 2, 3, 1, 2, 3].
|
||
- 'interleave': Repeat each prompt consecutively before moving to
|
||
the next. Example: [1, 2, 3] -> [1, 1, 2, 2, 3, 3].
|
||
|
||
Returns:
|
||
A list of repeated prompts in the specified order.
|
||
|
||
Raises:
|
||
ValueError: If an invalid mode is provided.
|
||
"""
|
||
write_resp(f"Repeat mode: {mode}\n")
|
||
if mode == "random":
|
||
repeated_prompts = prompts * repeat_count
|
||
random.shuffle(repeated_prompts)
|
||
return repeated_prompts
|
||
elif mode == "tile":
|
||
return prompts * repeat_count
|
||
elif mode == "interleave":
|
||
repeated_prompts = []
|
||
for prompt in prompts:
|
||
repeated_prompts.extend([prompt] * repeat_count)
|
||
return repeated_prompts
|
||
else:
|
||
raise ValueError(
|
||
f"Invalid mode: {mode}, only support 'random', 'tile', 'interleave'"
|
||
)
|
||
|
||
|
||
def add_cache_misses(prompts, hit_miss_ratio):
|
||
"""
|
||
Add cache misses to the prompts and return a boolean mask aligned with prompts.
|
||
"""
|
||
if hit_miss_ratio is None:
|
||
return prompts, [False] * len(prompts)
|
||
|
||
hit, miss = map(int, hit_miss_ratio.split(":", 1))
|
||
period = hit + miss
|
||
miss_mask = [False] * len(prompts)
|
||
|
||
for i in range(len(prompts)):
|
||
# every (hit+miss) window: first `hit` are hits, rest are misses
|
||
if period and (i % period) >= hit:
|
||
miss_mask[i] = True
|
||
prompts[i] = f"{random.randint(-10_000_000, 10_000_000)} {prompts[i]}"
|
||
|
||
return prompts, miss_mask
|
||
|
||
|
||
def relative_time(df, start_time):
|
||
"""
|
||
Relative time to the start of the benchmark.
|
||
"""
|
||
df["request_start"] = df["request_start"] - start_time
|
||
df["request_end"] = df["request_end"] - start_time
|
||
df["ttft_time"] = df["request_start"] + df["ttft"]
|
||
|
||
|
||
def visualize_results(warmup_df, benchmark_df):
|
||
def plot_bars(df, title, filename):
|
||
plt.figure(figsize=(12, 6))
|
||
|
||
if "is_miss" in df.columns:
|
||
is_miss = df["is_miss"]
|
||
else:
|
||
is_miss = pd.Series(False, index=df.index)
|
||
|
||
hits = df[~is_miss]
|
||
misses = df[is_miss]
|
||
|
||
# Prefill: dark blue (hit), dark orange (miss)
|
||
if not hits.empty:
|
||
plt.barh(
|
||
hits["prompt_id"],
|
||
hits["ttft_time"] - hits["request_start"],
|
||
left=hits["request_start"],
|
||
color="darkblue",
|
||
label="Loading", # prefill hits
|
||
)
|
||
if not misses.empty:
|
||
plt.barh(
|
||
misses["prompt_id"],
|
||
misses["ttft_time"] - misses["request_start"],
|
||
left=misses["request_start"],
|
||
color="darkorange",
|
||
label="Compute", # prefill misses
|
||
)
|
||
|
||
# Decode: light blue (hit), light orange (miss)
|
||
if not hits.empty:
|
||
plt.barh(
|
||
hits["prompt_id"],
|
||
hits["request_end"] - hits["ttft_time"],
|
||
left=hits["ttft_time"],
|
||
color="skyblue",
|
||
label="Decoding after loading",
|
||
)
|
||
if not misses.empty:
|
||
plt.barh(
|
||
misses["prompt_id"],
|
||
misses["request_end"] - misses["ttft_time"],
|
||
left=misses["ttft_time"],
|
||
color="pink",
|
||
label="Decoding after compute",
|
||
)
|
||
|
||
plt.xlabel("Time (s)")
|
||
plt.ylabel("Prompt ID")
|
||
plt.legend()
|
||
plt.tight_layout()
|
||
plt.savefig(filename)
|
||
plt.close()
|
||
|
||
plot_bars(warmup_df, "Warmup Round", "warmup_round.png")
|
||
plot_bars(benchmark_df, "Query Round", "query_round.png")
|
||
|
||
|
||
def trimmed_mean(series: pd.Series, trim_fraction: float) -> float:
|
||
"""
|
||
Exclude the smallest trim_fraction and largest trim_fraction and take mean
|
||
of the rest. If trim_fraction <= 0, returns normal mean.
|
||
"""
|
||
s = series.dropna()
|
||
if len(s) == 0:
|
||
return float("nan")
|
||
if trim_fraction <= 0:
|
||
return float(s.mean())
|
||
|
||
if not (0.0 <= trim_fraction < 0.5):
|
||
raise ValueError("--trim-fraction must be in [0, 0.5).")
|
||
|
||
s = s.sort_values()
|
||
n = len(s)
|
||
k = int(n * trim_fraction)
|
||
if n - 2 * k <= 0:
|
||
return float(s.mean())
|
||
return float(s.iloc[k : n - k].mean())
|
||
|
||
|
||
async def main(args):
|
||
random.seed(args.shuffle_seed)
|
||
|
||
# Create the OpenAI client
|
||
# No timeout: some benchmarks can take 4-5 minutes per request
|
||
base_url = get_url_from_args(args)
|
||
print("Using base URL:", base_url)
|
||
|
||
api_key = os.getenv("OPENAI_API_KEY", "sk-dummy")
|
||
|
||
client = AsyncOpenAI(
|
||
base_url=base_url,
|
||
api_key=api_key,
|
||
timeout=None,
|
||
)
|
||
model = args.model
|
||
if model == "auto":
|
||
print("Auto-selecting model...", end=" ")
|
||
models = (await client.models.list(),)
|
||
model = models[0].data[0].id
|
||
print(f"selected model: {model}")
|
||
|
||
pre_warmup_prompts = [str(i) + "xx" + " ".join(["hi"] * 1000) for i in range(5)]
|
||
|
||
await test_long_document_qa(
|
||
client=client,
|
||
model=model,
|
||
prompts=pre_warmup_prompts,
|
||
output_len=args.output_len,
|
||
max_inflight_requests=args.max_inflight_requests,
|
||
)
|
||
|
||
# Prepare the prompts:
|
||
# we append the document id at the beginning to avoid any of the document
|
||
# being the prefix of other documents
|
||
warmup_prompts = [
|
||
str(i) + " " + " ".join(["hi"] * args.document_length)
|
||
for i in range(args.num_documents)
|
||
]
|
||
|
||
prompts = repeat_prompts(warmup_prompts, args.repeat_count, mode=args.repeat_mode)
|
||
prompts, miss_mask = add_cache_misses(prompts, args.hit_miss_ratio)
|
||
|
||
write_resp("------warm up round------\n")
|
||
warmup_start_time = time.time()
|
||
warmup_request_stats = await test_long_document_qa(
|
||
client=client,
|
||
model=model,
|
||
prompts=warmup_prompts,
|
||
output_len=args.output_len,
|
||
max_inflight_requests=args.max_inflight_requests,
|
||
)
|
||
warmup_end_time = time.time()
|
||
write_resp("------query round------\n")
|
||
|
||
sleep_time_after_warmup = args.sleep_time_after_warmup
|
||
if sleep_time_after_warmup > 0:
|
||
write_resp(f"Sleeping for {sleep_time_after_warmup} seconds after warmup...\n")
|
||
time.sleep(sleep_time_after_warmup)
|
||
|
||
benchmark_start_time = time.time()
|
||
benchmark_request_stats = await test_long_document_qa(
|
||
client=client,
|
||
model=model,
|
||
prompts=prompts,
|
||
output_len=args.output_len,
|
||
max_inflight_requests=args.max_inflight_requests,
|
||
)
|
||
benchmark_end_time = time.time()
|
||
|
||
warmup_df = pd.DataFrame([stats.__dict__ for stats in warmup_request_stats])
|
||
relative_time(warmup_df, warmup_start_time)
|
||
warmup_df["is_miss"] = True
|
||
benchmark_df = pd.DataFrame([stats.__dict__ for stats in benchmark_request_stats])
|
||
benchmark_df["is_miss"] = miss_mask
|
||
relative_time(benchmark_df, benchmark_start_time)
|
||
|
||
warmup_df.to_csv("warmup_round.csv", index=False)
|
||
benchmark_df.to_csv("query_round.csv", index=False)
|
||
|
||
# Print results
|
||
warmup_mean_ttft = trimmed_mean(
|
||
warmup_df.query("successful == True")["ttft"], args.trim_fraction
|
||
)
|
||
query_mean_ttft = trimmed_mean(
|
||
benchmark_df.query("successful == True")["ttft"], args.trim_fraction
|
||
)
|
||
warmup_success_count = warmup_df.query("successful == True").shape[0]
|
||
query_success_count = benchmark_df.query("successful == True").shape[0]
|
||
CSI = "\x1b["
|
||
RESET = CSI + "0m"
|
||
print(f"Warmup round mean TTFT: {warmup_mean_ttft:.3f}s")
|
||
print(f"Warmup round time: {warmup_end_time - warmup_start_time:.3f}s")
|
||
print(f"Warmup round prompt count: {len(warmup_df)}")
|
||
print(f"Warmup round successful prompt count: {warmup_success_count}")
|
||
print(f"{CSI}36;1m\n=== BENCHMARK RESULTS ==={RESET}")
|
||
print(f"{CSI}32mQuery round mean TTFT: {query_mean_ttft:.3f}s{RESET}")
|
||
print(
|
||
f"{CSI}33mQuery round time: "
|
||
f"{benchmark_end_time - benchmark_start_time:.3f}s{RESET}"
|
||
)
|
||
print(f"{CSI}35mQuery round prompt count: {len(benchmark_df)}{RESET}")
|
||
print(f"{CSI}34mQuery round successful prompt count: {query_success_count}{RESET}")
|
||
|
||
if visualize:
|
||
visualize_results(warmup_df, benchmark_df)
|
||
|
||
if args.json_output:
|
||
query_duration = benchmark_end_time - benchmark_start_time
|
||
query_round_time_per_prompt = query_duration / len(benchmark_df)
|
||
warmup_duration = warmup_end_time - warmup_start_time
|
||
warmup_round_time_per_prompt = warmup_duration / len(warmup_df)
|
||
# Standard
|
||
import json
|
||
|
||
summary = {
|
||
"query_ttft_per_prompt": query_mean_ttft,
|
||
"query_round_time_per_prompt": query_round_time_per_prompt,
|
||
"warmup_round_time_per_prompt": warmup_round_time_per_prompt,
|
||
}
|
||
print(json.dumps(summary))
|
||
|
||
|
||
def create_argument_parser():
|
||
parser = argparse.ArgumentParser(
|
||
description="Benchmark the performance with or "
|
||
"without automatic prefix caching."
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--document-length",
|
||
type=int,
|
||
# Roughly the number of tokens for a system paper,
|
||
# excluding images
|
||
default=20000,
|
||
help="Length of each document in tokens.",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--num-documents",
|
||
type=int,
|
||
default=8,
|
||
help="Number of documents to generate for testing.",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--output-len",
|
||
type=int,
|
||
default=100,
|
||
help="Maximum number of tokens to generate for each prompt.",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--repeat-count",
|
||
type=int,
|
||
default=2,
|
||
help="Number of times to repeat each prompt",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--repeat-mode",
|
||
type=str,
|
||
default="random",
|
||
help="The mode to repeat prompts. The supported "
|
||
'modes are "random", "tile", and "interleave". '
|
||
"See repeat_prompts() in the source code for details.",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--shuffle-seed",
|
||
type=int,
|
||
default=0,
|
||
help='Random seed when the repeat mode is "random"',
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--host",
|
||
type=str,
|
||
default=None,
|
||
help="Host to query the vLLM server",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--port",
|
||
type=int,
|
||
default=None,
|
||
help="Port to query the vLLM server",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--base-url",
|
||
type=str,
|
||
default=None,
|
||
help="Base URL to query the LLM server",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--model",
|
||
type=str,
|
||
default="auto",
|
||
help="Model name, can be set to 'auto' if the "
|
||
"endpoint support openai api /models",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--max-inflight-requests",
|
||
type=int,
|
||
default=2,
|
||
help="Maximum number of concurrent inflight requests",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--sleep-time-after-warmup",
|
||
type=float,
|
||
default=0.0,
|
||
help="Sleep time after warm up iteration",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--output",
|
||
type=str,
|
||
default=None,
|
||
help="Filename to write all responses to; if omitted, writes to stdout.",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--completions",
|
||
action="store_true",
|
||
help="Use completions API instead of chat completions API",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--visualize",
|
||
action="store_true",
|
||
help="Visualize the results",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--hit-miss-ratio",
|
||
type=str,
|
||
default=None,
|
||
help=(
|
||
"In query round, control how many of the prompts will miss the cache."
|
||
"For example, 3:1 means every fourth repeated prompt will be randomized "
|
||
"to force a cache miss. 2:2 means 2 hits and 2 misses"
|
||
),
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--eos-token-id",
|
||
type=int,
|
||
default=None,
|
||
help=(
|
||
"EOS token id. we bias against this token id so we always "
|
||
"get the number of output tokens we specify"
|
||
),
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--json-output",
|
||
action="store_true",
|
||
help="Print benchmark summary as a single JSON line to stdout.",
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--trim-fraction",
|
||
type=float,
|
||
default=0.0,
|
||
help=(
|
||
"Exclude the smallest and largest fraction of successful samples "
|
||
"before averaging. Example: 0.1 drops bottom 10% and top 10%."
|
||
),
|
||
)
|
||
|
||
return parser
|
||
|
||
|
||
def validate_args(args):
|
||
# Verify port and base_url are exclusive
|
||
has_host_port = args.host is not None and args.port is not None
|
||
has_base_url = args.base_url is not None
|
||
if has_host_port and has_base_url:
|
||
raise ValueError("Cannot use --host/--port and --base-url together.")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = create_argument_parser()
|
||
args = parser.parse_args()
|
||
validate_args(args)
|
||
completions_mode = args.completions
|
||
visualize = args.visualize
|
||
if visualize:
|
||
# Third Party
|
||
import matplotlib.pyplot as plt
|
||
if args.eos_token_id is not None:
|
||
eos_token_id = args.eos_token_id
|
||
OUTPUT_FILE = args.output
|
||
asyncio.run(main(args))
|