285 lines
11 KiB
Python
285 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
import subprocess
|
|
import time
|
|
import requests
|
|
import json
|
|
import os
|
|
import sys
|
|
import argparse
|
|
from datetime import datetime
|
|
from typing import List, Tuple, Optional
|
|
|
|
# 配置常量
|
|
LOG_FILE = "llama_test_logs.csv"
|
|
CONTENT_FILE = "llama_content_logs.csv"
|
|
MAX_RETRIES = 10
|
|
RETRY_DELAY = 15
|
|
SERVER_PATH = "./build/bin/llama-server"
|
|
|
|
# 固定的测试prompt
|
|
FIXED_PROMPT = "who are you?"
|
|
|
|
def init_log_file():
|
|
"""初始化日志文件，写入CSV表头（如果文件不存在）"""
|
|
if not os.path.exists(LOG_FILE):
|
|
with open(LOG_FILE, "w") as f:
|
|
f.write("timestamp,model_path,case_path,case_name,prompt_tokens,completion_tokens,memory_usage,Decode_TPS,TTFT,Prefill_TPS,Power,result\n")
|
|
print(f"Initialized log file: {LOG_FILE}")
|
|
if not os.path.exists(CONTENT_FILE):
|
|
with open(CONTENT_FILE, "w") as f:
|
|
f.write("timestamp,model_path,case_path,case_name,prompt,content\n")
|
|
print(f"Initialized content file: {CONTENT_FILE}")
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description='Automatically test model (.gguf) files in case directories with fixed prompt.'
|
|
)
|
|
parser.add_argument(
|
|
"-c", "--case_path",
|
|
required=True,
|
|
help="Root directory to scan for case folders containing .gguf model files."
|
|
)
|
|
parser.add_argument(
|
|
"--port",
|
|
type=int,
|
|
default=8080,
|
|
help="llama-server port"
|
|
)
|
|
return parser.parse_args()
|
|
|
|
def discover_gguf_models(root_dir: str) -> List[Tuple[str, str, str]]:
|
|
"""
|
|
遍历根目录，查找所有包含.gguf文件的文件夹。
|
|
返回一个列表，元素为 (model_path, case_path, case_name) 元组。
|
|
case_path 是包含.gguf文件的文件夹路径，case_name 是文件夹名称。
|
|
"""
|
|
test_cases = []
|
|
|
|
if not os.path.isdir(root_dir):
|
|
print(f"Error: Provided root_dir '{root_dir}' is not a valid directory.")
|
|
return test_cases
|
|
|
|
print(f"Scanning root directory: {root_dir}")
|
|
|
|
for item in os.scandir(root_dir):
|
|
if item.is_dir():
|
|
case_path = item.path
|
|
case_name = item.name
|
|
gguf_files = []
|
|
|
|
# 遍历文件夹内所有文件，查找.gguf文件
|
|
for file_item in os.scandir(case_path):
|
|
if file_item.is_file() and file_item.name.endswith('.gguf'):
|
|
gguf_files.append(file_item.path)
|
|
|
|
if gguf_files:
|
|
# 如果有多个.gguf文件，可以选择第一个或让用户选择
|
|
# 这里我们选择第一个找到的.gguf文件
|
|
model_path = gguf_files[0]
|
|
if len(gguf_files) > 1:
|
|
print(f" Warning: Found {len(gguf_files)} .gguf files in '{case_name}', using: {os.path.basename(model_path)}")
|
|
test_cases.append((model_path, case_path, case_name))
|
|
print(f" Found: Case='{case_name}', Model='{os.path.basename(model_path)}'")
|
|
else:
|
|
print(f" Skipping folder '{case_name}': No .gguf files found.")
|
|
|
|
print(f"Total discovered test cases: {len(test_cases)}\n")
|
|
return test_cases
|
|
|
|
def start_server(model_path: str, case_path: str, port: int) -> Optional[subprocess.Popen]:
|
|
"""启动llama-server进程，返回subprocess.Popen对象。如果启动失败，返回None。"""
|
|
cmd = ["sudo", SERVER_PATH, "-m", str(model_path), "--case", str(case_path), "--port", str(port), "-np", str(1), "--verbose", "--no-warmup"]
|
|
print(f"Starting server with: Model='{model_path}', Case='{case_path}'")
|
|
print(f"Command: {' '.join(cmd)}")
|
|
try:
|
|
process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
time.sleep(1) # 稍等片刻让进程启动
|
|
return process
|
|
except Exception as e:
|
|
print(f"Failed to start server: {e}")
|
|
return None
|
|
|
|
def health_check(health_url: str) -> bool:
|
|
"""健康检查，重试直到服务就绪或超时"""
|
|
for i in range(1, MAX_RETRIES + 1):
|
|
try:
|
|
response = requests.get(health_url, timeout=15)
|
|
if response.status_code == 200:
|
|
print(f"✔ llama-server is ready (attempt {i}/{MAX_RETRIES})")
|
|
return True
|
|
except requests.exceptions.RequestException as e:
|
|
pass # 忽略连接错误
|
|
print(f" Waiting for service to be ready (attempt {i}/{MAX_RETRIES}, retry in {RETRY_DELAY}s)...")
|
|
time.sleep(RETRY_DELAY)
|
|
print("✘ llama-server startup timed out")
|
|
return False
|
|
|
|
def send_request(prompt: str, model_path: str, api_url: str) -> Optional[dict]:
|
|
"""发送对话请求，提取并返回性能数据。如果失败，返回None。"""
|
|
request_body = {
|
|
"model": model_path,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"stream": False
|
|
}
|
|
print(f" Sending request with prompt: '{prompt[:50]}...'" if len(prompt) > 50 else f" Sending request: '{prompt}'")
|
|
try:
|
|
response = requests.post(api_url, json=request_body, timeout=150)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
except requests.exceptions.Timeout:
|
|
print(f" ✗ Request timed out.")
|
|
return None
|
|
except requests.exceptions.RequestException as e:
|
|
print(f" Request failed: {e}")
|
|
return None
|
|
except json.JSONDecodeError as e:
|
|
print(f" JSON decode failed: {e}")
|
|
return None
|
|
|
|
# 提取数据
|
|
usage = data.get("usage", {})
|
|
timings = data.get("timings", {})
|
|
prompt_tokens = usage.get("prompt_tokens", 0)
|
|
completion_tokens = usage.get("completion_tokens", 0)
|
|
decode_tps = timings.get("predicted_per_second", 0.0)
|
|
predicted_ms = timings.get("predicted_ms", 0.0)
|
|
prompt_ms = timings.get("prompt_ms", 0.0)
|
|
ttft = predicted_ms + prompt_ms
|
|
prefill_tps = timings.get("prompt_per_second", 0.0)
|
|
message = data.get("choices", [{}])[0].get("message", {})
|
|
content = message.get("reasoning_content") or message.get("content", "")
|
|
|
|
return {
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"memory_usage": 0,
|
|
"decode_tps": decode_tps,
|
|
"ttft": ttft,
|
|
"prefill_tps": prefill_tps,
|
|
"Power": 0,
|
|
"prompt": prompt,
|
|
"content": content,
|
|
"model_path": model_path
|
|
}
|
|
|
|
def log_data(metrics: dict, case_path: str, case_name: str, result: str):
|
|
"""记录数据到CSV文件"""
|
|
if metrics is None:
|
|
print("No valid data, skipping logging.")
|
|
return
|
|
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
with open(LOG_FILE, "a") as f:
|
|
row = f"{timestamp},{metrics['model_path']},{case_path},{case_name},{metrics['prompt_tokens']},{metrics['completion_tokens']},{metrics['memory_usage']:.2f},{metrics['decode_tps']:.2f},{metrics['ttft']:.2f},{metrics['prefill_tps']:.2f},{metrics['Power']:.1f},{result}\n"
|
|
f.write(row)
|
|
with open(CONTENT_FILE, "a") as f:
|
|
# 对内容中的换行符和逗号进行处理，用空格替换换行，用占位符替换逗号以确保CSV格式正确
|
|
DELIMITER = '\t'
|
|
# safe_content = metrics['content'].replace('\n', ' ').replace(',', ';')
|
|
# safe_prompt = metrics['prompt'].replace('\n', ' ').replace(',', ';')
|
|
row = f"{timestamp},{metrics['model_path']}{DELIMITER}{case_path}{DELIMITER}{case_name}{DELIMITER}{metrics['prompt']}{DELIMITER}{metrics['content']}\n"
|
|
f.write(row)
|
|
|
|
def main():
|
|
"""主函数，执行完整流程"""
|
|
print("=== Starting Automated llama-server Test Suite (Fixed Prompt) ===")
|
|
|
|
# 1. 初始化日志文件
|
|
init_log_file()
|
|
|
|
# 2. 解析参数
|
|
args = parse_args()
|
|
|
|
# 设置固定prompt
|
|
#fixed_prompt = args.prompt
|
|
print(f"Using fixed prompt: '{FIXED_PROMPT}'")
|
|
|
|
# 发现所有包含.gguf文件的文件夹
|
|
test_cases = discover_gguf_models(args.case_path)
|
|
|
|
if not test_cases:
|
|
print("No folders with .gguf files found. Exiting.")
|
|
sys.exit(1)
|
|
|
|
overall_results = []
|
|
|
|
# 3. 遍历每个测试组合
|
|
for model_path, case_path, case_name in test_cases:
|
|
server_process = None
|
|
combination_result = "FAIL" # 默认失败
|
|
print(f"\n{'='*60}")
|
|
print(f"Testing Case:")
|
|
print(f" Case Name: {case_name}")
|
|
print(f" Case Path: {case_path}")
|
|
print(f" Model: {model_path}")
|
|
print(f"{'='*60}")
|
|
|
|
try:
|
|
# 启动服务器
|
|
server_process = start_server(model_path, case_path, args.port)
|
|
if not server_process:
|
|
raise RuntimeError(f"Failed to start server for model '{model_path}' and case '{case_path}'")
|
|
|
|
# 健康检查
|
|
health_url = f"http://localhost:{args.port}/health"
|
|
api_url = f"http://localhost:{args.port}/v1/chat/completions"
|
|
|
|
if not health_check(health_url):
|
|
raise RuntimeError("Health check failed. Server may not have started correctly.")
|
|
|
|
# 使用固定prompt发送请求
|
|
print(f"\n Sending test request with fixed prompt...")
|
|
|
|
metrics = send_request(FIXED_PROMPT, model_path, api_url)
|
|
|
|
if metrics:
|
|
# 记录请求的详细日志
|
|
log_data(metrics, case_path, case_name, "REQUEST_OK")
|
|
combination_result = "PASS"
|
|
print(f"\n ✔ Case '{case_name}' TEST PASSED: Request succeeded.")
|
|
else:
|
|
combination_result = "FAIL"
|
|
print(f"\n ✗ Case '{case_name}' TEST FAILED: Request failed.")
|
|
|
|
# 记录该组合的最终结果（汇总行）
|
|
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
# with open(LOG_FILE, "a") as f:
|
|
# summary_row = f"{timestamp},{model_path},{case_path},{case_name},,,,,,,,{combination_result}\n"
|
|
# f.write(summary_row)
|
|
|
|
overall_results.append((model_path, case_path, case_name, combination_result))
|
|
|
|
except Exception as e:
|
|
print(f"\n ✗ An error occurred during testing of case '{case_name}': {e}")
|
|
combination_result = "ERROR"
|
|
overall_results.append((model_path, case_path, case_name, f"ERROR: {e}"))
|
|
finally:
|
|
# 停止当前组合的服务器进程
|
|
if server_process and server_process.poll() is None:
|
|
print(f" Stopping server for case '{case_name}'...")
|
|
server_process.terminate()
|
|
try:
|
|
server_process.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
server_process.kill()
|
|
# 短暂等待，避免端口占用影响下一个测试
|
|
time.sleep(2)
|
|
|
|
# 4. 打印最终汇总报告
|
|
print(f"\n{'='*60}")
|
|
print("TEST SUITE SUMMARY:")
|
|
print(f"{'='*60}")
|
|
for model_path, case_path, case_name, result in overall_results:
|
|
status_symbol = "✔" if result == "PASS" else "✗"
|
|
print(f"{status_symbol} Case: {case_name:50} | Result: {result}")
|
|
|
|
# 判断整个测试套件是否全部通过
|
|
all_passed = all(result == "PASS" for _, _, _, result in overall_results)
|
|
if all_passed:
|
|
print("\n✔ All test cases PASSED.")
|
|
sys.exit(0)
|
|
else:
|
|
print("\n✗ Some test cases FAILED or encountered ERRORS.")
|
|
sys.exit(1)
|
|
|
|
if __name__ == "__main__":
|
|
main() |