Locust

简介:

是非常简单易用、分布式、python开发的压力测试工具。有图形化界面,支持将压测数据导出。

安装:

# pip3 安装locust
pip3  install locust
# 查看是否安装成功
locust -h
# 运行 Locust 分布在多个进程/机器库
pip3 install pyzmq
# webSocket 压测库
pip3 install websocket-client

使用示例:

# 默认使用 locusfile.py 文件
locust --host http://10.193.135.79:8880


# 使用自定义文件名
locust -f stress_from_excel.py \
  --host http://10.193.135.79:8880 \
  --users 100 \
  --spawn-rate 10


# 使用命令行模式(无需打开浏览器)
--headless --run-time 60s --csv report
from locust import HttpUser, task, between
from locust import events

import pandas as pd
import random

# ===== 读取 Excel 文件中用户 prompt =====

pc_corpus_file = r'/Users/zhaoweiguo/Downloads/todelete/intend/PC端语料.xlsx'
df = pd.read_excel(pc_corpus_file)  # 假设第一列是 "prompt"
prompts = df["query"].dropna().tolist()

# 可重复测试
random.seed(42)

total_num=0
total_fail=0


class ChatUser(HttpUser):
    wait_time = between(0.5, 1.5)  # 用户之间等待时间

    @task
    def chat_completion(self):
        global total_num, total_fail  # 声明要修改的全局变量

        prompt = random.choice(prompts)

        payload = {
            "model": "Qwen2/Qwen2.5",  # 与 curl 中一致
            "messages": [
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.7,
            "stream": False
        }

        headers = {
            "Authorization": "Bearer sk-xxxx",
            "Content-Type": "application/json",
            "User-Agent": "vscode-restclient"
        }

        response = self.client.post(
            "/v1/chat/completions", json=payload, headers=headers
        )
        total_num = total_num + 1
        
        if response.status_code != 200:
            total_fail = total_fail + 1
            print("响应状态码:", response.status_code)
            print("response:", response)
            print(f"请求失败,当前总请求数: {total_num}, 失败数: {total_fail}")
        
        # 随机打印,概率10%
        # if random.random() < 0.1:
        #     print("响应状态码:", response.status_code)
        #     print("响应内容:", response.text)


    @events.request.add_listener
    def on_request(request_type, name, response_time, response_length, response, context, exception, start_time, url, **kwargs):
        # print(f"----- [{request_type}] {name}, url: {url} - {response_time}, kwargs: {kwargs}")
        # print(f"----- response_time: {response_time}ms, start_time: {start_time}")
        if exception:
            print(f"----- [ERROR] {exception}")
        if response:
            if response.status_code != 200:
                print(f"----- [FAIL] {response.status_code}: {response.text}")