issue/16: 多请求多并发测试脚本

scripts/test_perf.py

+import asyncio
+import time
+from openai import AsyncOpenAI
+import argparse
+import random
+
+
+PROMPTS = [
+    "如果猫能写诗，它们会写些什么？",
+    "描述一个没有重力的世界。",
+    "如果地球停止自转，会发生什么？",
+    "假设你是一只会飞的鲸鱼，描述你的日常生活。",
+    "如果人类可以与植物沟通，世界会变成什么样？",
+    "描述一个由糖果构成的城市。",
+    "如果时间旅行成为可能，你最想去哪个时代？",
+    "想象一下，如果地球上只有蓝色，其他颜色都消失了。",
+    "如果动物能上网，它们会浏览什么网站？",
+    "描述一个没有声音的世界。",
+    "如果人类可以在水下呼吸，城市会如何变化？",
+    "想象一下，如果天空是绿色的，云是紫色的。",
+    "如果你能与任何历史人物共进晚餐，你会选择谁？",
+    "描述一个没有夜晚的星球。",
+    "如果地球上只有一种语言，世界会如何运作？",
+    "想象一下，如果所有的书都变成了音乐。",
+    "如果你可以变成任何一种动物，你会选择什么？",
+    "描述一个由机器人统治的未来世界。",
+    "如果你能与任何虚构角色成为朋友，你会选择谁？",
+    "想象一下，如果每个人都能读懂他人的思想。"
+]
+
+
+async def benchmark_user(client, semaphore, queue, results, model):
+    while True:
+        async with semaphore:
+            task_id = await queue.get()
+            if task_id is None:
+                queue.task_done()
+                break
+            print(f"🚀 Sending request #{task_id}")
+            start_time = time.time()
+            content =  random.choice(PROMPTS)
+            stream = await client.chat.completions.create(
+                model=model,
+                messages=[{"role": "user", "content": content}],
+                stream=True
+            )
+
+            first_token_time = None
+            total_tokens = 0
+            async for chunk in stream:
+                if first_token_time is None:
+                    first_token_time = time.time()
+                if chunk.choices[0].delta.content:
+                    # print(chunk.choices[0].delta.content, end="", flush=True)
+                    total_tokens += 1
+                if chunk.choices[0].finish_reason is not None:
+                    break
+
+            end_time = time.time()
+            elapsed_time = end_time - start_time
+            ttft = first_token_time - start_time if first_token_time else None
+            tokens_per_second = total_tokens / elapsed_time if elapsed_time > 0 else 0
+            ms_per_token = (elapsed_time / total_tokens * 1000) if total_tokens > 0 else None
+
+            results.append((total_tokens, elapsed_time, tokens_per_second, ttft, ms_per_token))
+            queue.task_done()
+
+
+async def run_benchmark(num_requests, concurrency,  llm_url, model):
+    client = AsyncOpenAI(base_url=llm_url, api_key="default")
+    semaphore = asyncio.Semaphore(concurrency)
+    queue = asyncio.Queue()
+    results = []
+
+    for i in range(num_requests):
+        await queue.put(i)
+    
+    for _ in range(concurrency):
+        await queue.put(None)
+
+
+    users = [asyncio.create_task(benchmark_user(client, semaphore, queue, results, model)) for _ in range(concurrency)]
+
+    start_time = time.time()
+    
+    await queue.join()
+    await asyncio.gather(*users)
+    end_time = time.time()
+
+    # Calculate metrics
+    total_elapsed_time = end_time - start_time
+    tokens_list = [r[0] for r in results if r and r[0] is not None]
+    latencies = [r[1] for r in results if r and r[1] is not None]
+    tokens_per_second_list = [r[2] for r in results if r and r[2] is not None]
+    ttft_list = [r[3] for r in results if r and r[3] is not None]
+    ms_per_token_list = [r[4] for r in results if r and r[4] is not None]
+
+    successful_requests = len(results)
+    requests_per_second = successful_requests / total_elapsed_time if total_elapsed_time > 0 else 0
+    avg_latency = sum(latencies) / len(latencies) if latencies else 0
+    avg_tokens_per_second = sum(tokens_per_second_list) / len(tokens_per_second_list) if tokens_per_second_list else 0
+    avg_ttft = sum(ttft_list) / len(ttft_list) if ttft_list else 0
+    avg_ms_per_token = sum(ms_per_token_list) / len(ms_per_token_list) if ms_per_token_list else None
+
+    width_label = 18
+    sep = "-" * 50
+
+    print(f"\n=== 📊 性能指标汇总 ({model}) ===")
+    print(sep)
+    print(f"{'并发数':<{width_label}}: {concurrency}")
+    print(f"{'请求总数':<{width_label}}: {num_requests}")
+    print(f"{'成功请求数':<{width_label}}: {successful_requests}")
+    print(f"{'总耗时':<{width_label}}: {total_elapsed_time:.2f} s")
+    print(f"{'总输出token数':<{width_label}}: {sum(tokens_list)}")
+    print(f"{'请求速率 (RPS)':<{width_label}}: {requests_per_second:.2f} requests/s")
+    print(sep)
+    print(f"{'Average latency':<{width_label}}: {avg_latency:.2f} s")
+    print(f"{'Average TTFT':<{width_label}}: {avg_ttft:.2f} s")
+    print(f"{'Avg time per token':<{width_label}}: {avg_ms_per_token:.2f} ms/token")
+    print(f"{'Avg Token generation speed':<{width_label}}: {avg_tokens_per_second:.2f} tokens/s")
+    print(sep)
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_requests", type=int, required=True)
+    parser.add_argument("--concurrency", type=int, required=True)
+    parser.add_argument("--api_url", type=str, required=True)
+    parser.add_argument("--model", type=str, default="FM9G-7B")
+    args = parser.parse_args()
+
+    asyncio.run(run_benchmark(
+        args.num_requests, 
+        args.concurrency, 
+        args.api_url, 
+        args.model,
+    ))