# SPDX-License-Identifier: Apache-2.0 # Standard from dataclasses import dataclass import argparse import asyncio import logging import random import time # Third Party from transformers import AutoTokenizer from utils import ( AsyncLoopWrapper, PromptBuildMethodType, build_rag_prompt, compute_f1, compute_rl, init_logger, load_dataset, ) import openai import pandas as pd logger = init_logger(__name__, logging.INFO) system_prompt_set = { PromptBuildMethodType.QA: "You will be asked a question after reading several passages. " # noqa: E501 "Please directly answer the question based on the given passages. " "Do NOT repeat the question. " "The answer should be within 5 words..\nPassages:\n", PromptBuildMethodType.FEW_SHOT: "Summarize the dialogue into a few short sentences. " # noqa: E501 "The following are some examples.\n\n", } query_prompt_set = { PromptBuildMethodType.QA: "\n\nAnswer the question directly based on the given passages." # noqa: E501 " Do NOT repeat the question. " "The answer should be within 5 words. \nQuestion:", PromptBuildMethodType.FEW_SHOT: "", } @dataclass class WorkloadConfig: # Overall QPS qps: float # Model name model: str # Tokenizer name tokenizer: str # Dataset. dataset: str # Start index of the workload start_index: int # End index of the workload end_index: int # Random shuffle. shuffle: bool # System prompt. system_prompt: str # Separator. separator: str # Query prompt. query_prompt: str # Prompt build method. prompt_build_method: PromptBuildMethodType # Max tokens for each generation. max_tokens: int @dataclass class Response: request_id: int body: str ttft: float generation_time: float prompt_tokens: int generation_tokens: int launch_time: float finish_time: float def parse_arguments(): parser = argparse.ArgumentParser(description="Parse RAG benchmark configurations.") parser.add_argument("--qps", type=float, required=True, help="Overall QPS") parser.add_argument("--model", type=str, required=True, help="Model name") parser.add_argument("--tokenizer", type=str, default="", help="Tokenizer name") parser.add_argument("--dataset", type=str, required=True, help="The dataset path") parser.add_argument( "--start-index", type=int, default=0, help="Start index of the workload" ) parser.add_argument( "--end-index", type=int, default=-1, help="End index of the workload" ) parser.add_argument("--shuffle", action="store_true", help="Random shuffle") parser.add_argument("--system-prompt", type=str, default="", help="System prompt") parser.add_argument("--separator", type=str, default="", help="Separator") parser.add_argument("--query-prompt", type=str, default="", help="Query prompt") parser.add_argument( "--prompt-build-method", type=str, required=True, help="Prompt build method", ) parser.add_argument( "--base-url", type=str, required=True, help="Base URL of the serving engine endpoint", ) parser.add_argument( "--api-key", type=str, default="EMPTY", help="API key of the serving engine endpoint", ) parser.add_argument( "--output", type=str, default="summary.csv", help="The output file name (ended with csv or txt) for the summary csv and txt", ) parser.add_argument( "--warmup", action="store_true", help="Whether to enable warmup" ) parser.add_argument( "--time", type=int, default=None, help="The total running time in seconds", ) parser.add_argument( "--verbose", action="store_true", help="Whether to enable verbose logging", ) parser.add_argument( "--max-tokens", type=int, default=32, help="Max tokens for each generation", ) parser.add_argument( "--step-interval", type=float, default=0.02, help="Step interval" ) args = parser.parse_args() return args class RequestExecutor: def __init__( self, base_url: str, api_key: str, prompt_build_method: PromptBuildMethodType, model: str, ): self.client = openai.AsyncOpenAI(api_key=api_key, base_url=base_url) self.model = model self.loop = AsyncLoopWrapper.GetOrStartLoop() self.prompt_build_method = prompt_build_method async def _async_launch_request(self, request_id, prompt, max_tokens): start_time = time.time() first_token_time = None words = "" response = None if self.prompt_build_method == PromptBuildMethodType.QA: messages = [{"role": "user", "content": prompt}] response = await self.client.chat.completions.create( messages=messages, model=self.model, temperature=0, stream=True, max_tokens=max_tokens, stream_options={"include_usage": True}, ) elif self.prompt_build_method == PromptBuildMethodType.FEW_SHOT: response = await self.client.completions.create( prompt=prompt, model=self.model, temperature=0, stream=True, max_tokens=max_tokens, stream_options={"include_usage": True}, ) else: raise ValueError(f"Invalid prompt build method {self.prompt_build_method}") async for tok in response: if not tok.choices: continue chunk_message = tok.choices[0].delta.content if chunk_message is not None: if first_token_time is None and chunk_message != "": first_token_time = time.time() words += chunk_message tokens_out = tok.usage.completion_tokens tokens_prefill = tok.usage.prompt_tokens finish_time = time.time() if first_token_time is None: first_token_time = finish_time return Response( request_id=request_id, body=words, ttft=first_token_time - start_time, generation_time=finish_time - first_token_time, prompt_tokens=tokens_prefill, generation_tokens=tokens_out, launch_time=start_time, finish_time=finish_time, ) def launch_request(self, request_id: int, prompt, max_tokens, finish_callback): """ finish_callback: Callable[[Response], None] """ real_callback = lambda x: finish_callback(x.result()) future = asyncio.run_coroutine_threadsafe( self._async_launch_request(request_id, prompt, max_tokens), self.loop, ) future.add_done_callback(real_callback) def warmup_engine(executor: RequestExecutor): logger.info("Warming up the engine") for i in range(10): prompt = f"WARMUP: Hi, I'm user {i}. Here are some text: {'hi ' * 100}." executor.launch_request(-1, prompt, 100, lambda x: None) AsyncLoopWrapper.WaitLoop() logger.info("Warm up finished.") class RAGManager: def __init__(self, workload_config: WorkloadConfig): self.workload_config = workload_config eval_dataset = load_dataset(workload_config.dataset) start_index = workload_config.start_index end_index = workload_config.end_index if end_index < 0: end_index = len(eval_dataset) eval_dataset = eval_dataset[start_index:end_index] if workload_config.shuffle: random.shuffle(eval_dataset) self._prompts = [] self._answers = [] self._build_method = workload_config.prompt_build_method self._generated_text: list[str | None] = [] self._generation_time: list[float | None] = [] self._prefill_tok_cnt: list[int | None] = [] self._generation_tok_cnt: list[int | None] = [] self._ttft: list[float | None] = [] self._tpot: list[float | None] = [] for ex in eval_dataset: prompt, _ = build_rag_prompt( workload_config.system_prompt, ex, workload_config.query_prompt, workload_config.separator, workload_config.prompt_build_method, ) self._prompts.append(prompt) self._answers.append(ex["answers"]) self._generated_text.append(None) self._generation_time.append(None) self._prefill_tok_cnt.append(None) self._generation_tok_cnt.append(None) self._ttft.append(None) self._tpot.append(None) self._tokenizer = AutoTokenizer.from_pretrained(workload_config.tokenizer) self._last_request_time = -1.0 self._last_request_index = 0 assert workload_config.qps > 0 self._gap = 1.0 / workload_config.qps self._max_tokens = workload_config.max_tokens def _update_result(self, response: Response): self._generated_text[response.request_id] = response.body self._ttft[response.request_id] = response.ttft self._tpot[response.request_id] = ( response.generation_time / response.generation_tokens ) self._generation_time[response.request_id] = response.generation_time self._prefill_tok_cnt[response.request_id] = response.prompt_tokens self._generation_tok_cnt[response.request_id] = response.generation_tokens def step(self, timestamp: float, executor: RequestExecutor) -> bool: if self._last_request_index >= len(self._prompts): return False if ( self._last_request_time < 0 or timestamp >= self._last_request_time + self._gap ): prompt = self._prompts[self._last_request_index] request_id = self._last_request_index self._last_request_time = timestamp self._last_request_index += 1 executor.launch_request( request_id, prompt, self._max_tokens, self._update_result ) return True def summary(self, start_time: float, end_time: float) -> pd.DataFrame: cnt = len(self._ttft) assert cnt > 0 avg_ttft = sum(ttft for ttft in self._ttft if ttft is not None) / cnt avg_tpot = sum(tpot for tpot in self._tpot if tpot is not None) / cnt # Create a dataframe quality = [] for i in range(cnt): if self._build_method == PromptBuildMethodType.QA: quality.append( max( [ compute_f1(self._generated_text[i], answer, self._tokenizer) for answer in self._answers[i] ] ) ) elif self._build_method == PromptBuildMethodType.FEW_SHOT: quality.append( max( [ compute_rl(self._generated_text[i], answer) for answer in self._answers[i] ] ) ) else: raise ValueError(f"Invalid prompt build method {self._build_method}") avg_quality = sum(quality) / cnt df = pd.DataFrame( { "quality": quality, "ttft": self._ttft, "tpot": self._tpot, "generation_time": self._generation_time, "prefill_token_cnt": self._prefill_tok_cnt, "generation_token_cnt": self._generation_tok_cnt, } ) total_time = end_time - start_time thput = cnt / total_time logger.info( f"Summary: {cnt} requests, average_ttft={avg_ttft} (second)\n" f" average_tpot={avg_tpot} (second)\n" f"throughput={thput} (req/s)\n" f"average_quality={avg_quality}\n" ) return df def run_rag(args): build_prompt_method_str = args.prompt_build_method.upper() build_prompt_method = None if build_prompt_method_str == "QA": build_prompt_method = PromptBuildMethodType.QA elif build_prompt_method_str == "FEW_SHOT": build_prompt_method = PromptBuildMethodType.FEW_SHOT else: raise ValueError(f"Invalid prompt build method {build_prompt_method_str}") workload_config = WorkloadConfig( qps=args.qps, model=args.model, tokenizer=args.tokenizer, dataset=args.dataset, start_index=args.start_index, end_index=args.end_index, shuffle=args.shuffle, system_prompt=args.system_prompt, separator=args.separator, query_prompt=args.query_prompt, prompt_build_method=build_prompt_method, max_tokens=args.max_tokens, ) executor = RequestExecutor( base_url=args.base_url, api_key=args.api_key, prompt_build_method=build_prompt_method, model=args.model, ) if args.warmup: warmup_engine(executor) manager = RAGManager(workload_config) step_interval = args.step_interval num_steps = 0 start_time = time.time() try: while True: num_steps += 1 effective = manager.step(time.time(), executor) if not effective: break time.sleep(step_interval) if args.time is not None and time.time() - start_time > args.time: break except KeyboardInterrupt: logger.info("Interrupted, waiting for the final result") AsyncLoopWrapper.StopLoop() logger.info(f"Finished benchmarking, dumping summary to {args.output}") summary = manager.summary(start_time, time.time()) summary.to_csv(args.output, index=False) def main(): args = parse_arguments() build_prompt_method_str = args.prompt_build_method.upper() build_prompt_method = None if build_prompt_method_str == "QA": build_prompt_method = PromptBuildMethodType.QA elif build_prompt_method_str == "FEW_SHOT": build_prompt_method = PromptBuildMethodType.FEW_SHOT else: raise ValueError(f"Invalid prompt build method {build_prompt_method_str}") if len(args.system_prompt) == 0: args.system_prompt = system_prompt_set[build_prompt_method] if len(args.query_prompt) == 0: args.query_prompt = query_prompt_set[build_prompt_method] if len(args.tokenizer) == 0: args.tokenizer = args.model args.system_prompt = args.system_prompt.encode().decode("unicode_escape") args.query_prompt = args.query_prompt.encode().decode("unicode_escape") if args.verbose: global logger logger = init_logger(__name__, logging.DEBUG) run_rag(args) if __name__ == "__main__": main()