serving.py 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214
  1. import argparse
  2. import asyncio
  3. import json
  4. import random
  5. import time
  6. from typing import AsyncGenerator, List, Tuple
  7. import aiohttp
  8. import numpy as np
  9. from transformers import PreTrainedTokenizerBase
  10. from aphrodite.transformers_utils.tokenizer import get_tokenizer
  11. # (prompt len, output len, latency)
  12. REQUEST_LATENCY: List[Tuple[int, int, float]] = []
  13. def sample_requests(
  14. dataset_path: str,
  15. num_requests: int,
  16. tokenizer: PreTrainedTokenizerBase,
  17. ) -> List[Tuple[str, int, int]]:
  18. # Load the dataset.
  19. with open(dataset_path) as f:
  20. dataset = json.load(f)
  21. # Filter out the conversations with less than 2 turns.
  22. dataset = [
  23. data for data in dataset
  24. if len(data["conversations"]) >= 2
  25. ]
  26. # Only keep the first two turns of each conversation.
  27. dataset = [
  28. (data["conversations"][0]["value"], data["conversations"][1]["value"])
  29. for data in dataset
  30. ]
  31. # Tokenize the prompts and completions.
  32. prompts = [prompt for prompt, _ in dataset]
  33. prompt_token_ids = tokenizer(prompts).input_ids
  34. completions = [completion for _, completion in dataset]
  35. completion_token_ids = tokenizer(completions).input_ids
  36. tokenized_dataset = []
  37. for i in range(len(dataset)):
  38. output_len = len(completion_token_ids[i])
  39. tokenized_dataset.append((prompts[i], prompt_token_ids[i], output_len))
  40. # Filter out too long sequences.
  41. filtered_dataset: List[Tuple[str, int, int]] = []
  42. for prompt, prompt_token_ids, output_len in tokenized_dataset:
  43. prompt_len = len(prompt_token_ids)
  44. if prompt_len < 4 or output_len < 4:
  45. # Prune too short sequences.
  46. continue
  47. if prompt_len > 1024 or prompt_len + output_len > 2048:
  48. # Prune too long sequences.
  49. continue
  50. filtered_dataset.append((prompt, prompt_len, output_len))
  51. # Sample the requests.
  52. sampled_requests = random.sample(filtered_dataset, num_requests)
  53. return sampled_requests
  54. async def get_request(
  55. input_requests: List[Tuple[str, int, int]],
  56. request_rate: float,
  57. ) -> AsyncGenerator[Tuple[str, int, int], None]:
  58. input_requests = iter(input_requests)
  59. for request in input_requests:
  60. yield request
  61. if request_rate == float("inf"):
  62. # If the request rate is infinity, then we don't need to wait.
  63. continue
  64. # Sample the request interval from the exponential distribution.
  65. interval = np.random.exponential(1.0 / request_rate)
  66. # The next request will be sent after the interval.
  67. await asyncio.sleep(interval)
  68. async def send_request(
  69. backend: str,
  70. api_url: str,
  71. prompt: str,
  72. prompt_len: int,
  73. output_len: int,
  74. best_of: int,
  75. use_beam_search: bool,
  76. ) -> None:
  77. request_start_time = time.time()
  78. headers = {"User-Agent": "Benchmark Client"}
  79. if backend == "aphrodite":
  80. pload = {
  81. "prompt": prompt,
  82. "n": 1,
  83. "best_of": best_of,
  84. "use_beam_search": use_beam_search,
  85. "temperature": 0.0 if use_beam_search else 1.0,
  86. "top_p": 1.0,
  87. "max_tokens": output_len,
  88. "ignore_eos": True,
  89. "stream": False,
  90. }
  91. elif backend == "tgi":
  92. assert not use_beam_search
  93. params = {
  94. "best_of": best_of,
  95. "max_new_tokens": output_len,
  96. "do_sample": True,
  97. }
  98. pload = {
  99. "inputs": prompt,
  100. "parameters": params,
  101. }
  102. else:
  103. raise ValueError(f"Unknown backend: {backend}")
  104. timeout = aiohttp.ClientTimeout(total=3 * 3600)
  105. async with aiohttp.ClientSession(timeout=timeout) as session:
  106. while True:
  107. async with session.post(api_url, headers=headers, json=pload) as response:
  108. chunks = []
  109. async for chunk, _ in response.content.iter_chunks():
  110. chunks.append(chunk)
  111. output = b"".join(chunks).decode("utf-8")
  112. output = json.loads(output)
  113. # Re-send the request if it failed.
  114. if "error" not in output:
  115. break
  116. request_end_time = time.time()
  117. request_latency = request_end_time - request_start_time
  118. REQUEST_LATENCY.append((prompt_len, output_len, request_latency))
  119. async def benchmark(
  120. backend: str,
  121. api_url: str,
  122. input_requests: List[Tuple[str, int, int]],
  123. best_of: int,
  124. use_beam_search: bool,
  125. request_rate: float,
  126. ) -> None:
  127. tasks: List[asyncio.Task] = []
  128. async for request in get_request(input_requests, request_rate):
  129. prompt, prompt_len, output_len = request
  130. task = asyncio.create_task(send_request(backend, api_url, prompt,
  131. prompt_len, output_len,
  132. best_of, use_beam_search))
  133. tasks.append(task)
  134. await asyncio.gather(*tasks)
  135. def main(args: argparse.Namespace):
  136. print(args)
  137. random.seed(args.seed)
  138. np.random.seed(args.seed)
  139. api_url = f"http://{args.host}:{args.port}/api/v1/generate"
  140. tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code)
  141. input_requests = sample_requests(args.dataset, args.num_prompts, tokenizer)
  142. benchmark_start_time = time.time()
  143. asyncio.run(benchmark(args.backend, api_url, input_requests, args.best_of,
  144. args.use_beam_search, args.request_rate))
  145. benchmark_end_time = time.time()
  146. benchmark_time = benchmark_end_time - benchmark_start_time
  147. print(f"Total time: {benchmark_time:.2f} s")
  148. print(f"Throughput: {args.num_prompts / benchmark_time:.2f} requests/s")
  149. # Compute the latency statistics.
  150. avg_latency = np.mean([latency for _, _, latency in REQUEST_LATENCY])
  151. print(f"Average latency: {avg_latency:.2f} s")
  152. avg_per_token_latency = np.mean([
  153. latency / (prompt_len + output_len)
  154. for prompt_len, output_len, latency in REQUEST_LATENCY
  155. ])
  156. print(f"Average latency per token: {avg_per_token_latency:.2f} s")
  157. avg_per_output_token_latency = np.mean([
  158. latency / output_len
  159. for _, output_len, latency in REQUEST_LATENCY
  160. ])
  161. print("Average latency per output token: "
  162. f"{avg_per_output_token_latency:.2f} s")
  163. if __name__ == "__main__":
  164. parser = argparse.ArgumentParser(
  165. description="Benchmark the online serving throughput.")
  166. parser.add_argument("--backend", type=str, default="aphrodite",
  167. choices=["aphrodite", "tgi"])
  168. parser.add_argument("--host", type=str, default="localhost")
  169. parser.add_argument("--port", type=int, default=8000)
  170. parser.add_argument("--dataset", type=str, required=True,
  171. help="Path to the dataset.")
  172. parser.add_argument("--tokenizer", type=str, required=True,
  173. help="Name or path of the tokenizer.")
  174. parser.add_argument("--best-of", type=int, default=1,
  175. help="Generates `best_of` sequences per prompt and "
  176. "returns the best one.")
  177. parser.add_argument("--use-beam-search", action="store_true")
  178. parser.add_argument("--num-prompts", type=int, default=1000,
  179. help="Number of prompts to process.")
  180. parser.add_argument("--request-rate", type=float, default=float("inf"),
  181. help="Number of requests per second. If this is inf, "
  182. "then all the requests are sent at time 0. "
  183. "Otherwise, we use Poisson process to synthesize "
  184. "the request arrival times.")
  185. parser.add_argument("--seed", type=int, default=0)
  186. parser.add_argument('--trust-remote-code', action='store_true',
  187. help='trust remote code from huggingface')
  188. args = parser.parse_args()
  189. main(args)