Add cache hit rate UT (#18566)
This commit is contained in:
Liangsheng Yin
@@ -36,20 +36,18 @@ class ContextWorkloadGenerator(WorkloadGenerator):
|
||||
init_requests = []
|
||||
for i in range(num_requests):
|
||||
context_id = self.dataset["queries"][i]["context"]
|
||||
init_requests.append(
|
||||
(
|
||||
i,
|
||||
gen_payload(
|
||||
self.dataset["contexts"][context_id]
|
||||
+ self.dataset["queries"][i]["question"],
|
||||
len(
|
||||
self.tokenizer(
|
||||
self.dataset["queries"][i]["reference_answer"]
|
||||
)["input_ids"]
|
||||
),
|
||||
),
|
||||
)
|
||||
# Tokenize the context + question to get input_ids
|
||||
prompt_text = (
|
||||
self.dataset["contexts"][context_id]
|
||||
+ self.dataset["queries"][i]["question"]
|
||||
)
|
||||
input_ids = self.tokenizer.encode(prompt_text)
|
||||
output_len = len(
|
||||
self.tokenizer(self.dataset["queries"][i]["reference_answer"])[
|
||||
"input_ids"
|
||||
]
|
||||
)
|
||||
init_requests.append((i, gen_payload(input_ids, output_len)))
|
||||
self.ready_queue = ReadyQueue(init_requests=init_requests)
|
||||
|
||||
self.response_queue = queue.Queue()
|
||||
|
||||
@@ -6,21 +6,13 @@ import random
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
import requests
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
from sglang.bench_serving import (
|
||||
RequestFuncOutput,
|
||||
get_tokenizer,
|
||||
remove_prefix,
|
||||
sample_random_requests,
|
||||
)
|
||||
|
||||
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60)
|
||||
from sglang.bench_serving import get_tokenizer, sample_random_requests
|
||||
from sglang.test.kits.cache_hit_kit import async_request_sglang_generate, gen_payload
|
||||
|
||||
|
||||
def parse_args():
|
||||
@@ -143,95 +135,6 @@ def parse_args():
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
async def async_request_sglang_generate(
|
||||
payload,
|
||||
url,
|
||||
pbar: Optional[tqdm] = None,
|
||||
):
|
||||
"""
|
||||
Sends a streaming request to the server. Gathers text token-by-token.
|
||||
"""
|
||||
async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
|
||||
headers = {}
|
||||
generated_text = ""
|
||||
ttft = 0.0
|
||||
st = time.perf_counter()
|
||||
most_recent_timestamp = st
|
||||
output = RequestFuncOutput()
|
||||
|
||||
try:
|
||||
async with session.post(url=url, json=payload, headers=headers) as response:
|
||||
if response.status == 200:
|
||||
prompt_tokens = 0
|
||||
cached_tokens = 0
|
||||
async for chunk_bytes in response.content:
|
||||
chunk_bytes = chunk_bytes.strip()
|
||||
if not chunk_bytes:
|
||||
continue
|
||||
|
||||
chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
|
||||
latency = time.perf_counter() - st
|
||||
if chunk == "[DONE]":
|
||||
pass
|
||||
else:
|
||||
data = json.loads(chunk)
|
||||
|
||||
if data["text"]:
|
||||
timestamp = time.perf_counter()
|
||||
# First token
|
||||
if ttft == 0.0:
|
||||
ttft = time.perf_counter() - st
|
||||
output.ttft = ttft
|
||||
prompt_tokens = (data.get("meta_info") or {}).get(
|
||||
"prompt_tokens", 0
|
||||
)
|
||||
cached_tokens = (data.get("meta_info") or {}).get(
|
||||
"cached_tokens", 0
|
||||
)
|
||||
|
||||
# Decoding phase
|
||||
else:
|
||||
output.itl.append(timestamp - most_recent_timestamp)
|
||||
|
||||
most_recent_timestamp = timestamp
|
||||
generated_text = data["text"]
|
||||
|
||||
output.generated_text = generated_text
|
||||
output.success = True
|
||||
output.latency = latency
|
||||
output.prompt_len = prompt_tokens
|
||||
output.cached_tokens = cached_tokens
|
||||
output.generated_len = len(output.itl) + 1
|
||||
else:
|
||||
output.error = response.reason or ""
|
||||
output.success = False
|
||||
except Exception as e:
|
||||
output.success = False
|
||||
output.error = str(e)
|
||||
print(f"Request failed: {e}")
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
|
||||
|
||||
def gen_payload(prompt, output_len, lora_path=""):
|
||||
payload = {
|
||||
"text": prompt,
|
||||
"sampling_params": {
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": output_len,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
"lora_path": lora_path,
|
||||
"return_logprob": False,
|
||||
"logprob_start_len": -1,
|
||||
}
|
||||
return payload
|
||||
|
||||
|
||||
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}
|
||||
@@ -286,6 +189,7 @@ class WorkloadGenerator:
|
||||
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,
|
||||
@@ -294,8 +198,10 @@ class WorkloadGenerator:
|
||||
tokenizer=self.tokenizer,
|
||||
dataset_path=args.dataset_path,
|
||||
random_sample=not args.disable_random_sample,
|
||||
return_text=False,
|
||||
)
|
||||
self.candidate_inputs = [i.prompt for i in self.candidate_inputs]
|
||||
# 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
|
||||
@@ -310,6 +216,7 @@ class WorkloadGenerator:
|
||||
tokenizer=self.tokenizer,
|
||||
dataset_path=args.dataset_path,
|
||||
random_sample=not args.disable_random_sample,
|
||||
return_text=False,
|
||||
)
|
||||
|
||||
init_requests = [
|
||||
@@ -321,8 +228,9 @@ class WorkloadGenerator:
|
||||
)
|
||||
for i in range(args.num_clients)
|
||||
]
|
||||
# history now stores List[int] (token ids) for each client
|
||||
self.client_records = {
|
||||
i: {"round": 0, "history": init_requests[i][1]["text"]}
|
||||
i: {"round": 0, "history": list(self.candidate_inputs[i])}
|
||||
for i in range(args.num_clients)
|
||||
}
|
||||
self.ready_queue = ReadyQueue(
|
||||
@@ -408,7 +316,8 @@ class WorkloadGenerator:
|
||||
) # Block until response is available
|
||||
if not response.success:
|
||||
raise ValueError(f"Request failed with error: {response.error}")
|
||||
self.client_records[client_id]["history"] += response.generated_text
|
||||
# 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)
|
||||
@@ -435,10 +344,9 @@ class WorkloadGenerator:
|
||||
self.completed_requests += 1
|
||||
|
||||
if self.client_records[client_id]["round"] < self.num_rounds:
|
||||
# append new request to client's history
|
||||
self.client_records[client_id][
|
||||
"history"
|
||||
] += self.sub_question_inputs.pop().prompt
|
||||
# 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(
|
||||
|
||||
Reference in New Issue
Block a user