data: publish complete Calculet NPU research archive
This commit is contained in:
@@ -0,0 +1,339 @@
|
||||
#!/usr/bin/env python3
|
||||
import subprocess
|
||||
import time
|
||||
import requests
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
from datetime import datetime
|
||||
import csv
|
||||
import random
|
||||
|
||||
# 配置常量
|
||||
LOG_FILE = "llama_perf_logs.csv"
|
||||
CONTENT_FILE = "llama_content_logs.csv"
|
||||
MAX_RETRIES = 10
|
||||
RETRY_DELAY = 15
|
||||
|
||||
def init_log_file():
|
||||
"""初始化日志文件，写入CSV表头（如果文件不存在）"""
|
||||
if not os.path.exists(LOG_FILE):
|
||||
with open(LOG_FILE, "w") as f:
|
||||
f.write("timestamp,model,prompt_tokens,completion_tokens,memory_usage,Decode_TPS,TTFT,Prefill_TPS,Power\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,prompt,content\n")
|
||||
print(f"Initialized content file: {CONTENT_FILE}")
|
||||
|
||||
def load_test_cases(test_file):
|
||||
"""从JSON文件加载测试用例"""
|
||||
try:
|
||||
with open(test_file, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
meta = data.get("meta_info", {})
|
||||
test_cases = data.get("test_cases", [])
|
||||
|
||||
print(f"filename: {test_file}")
|
||||
print(f"version: {meta.get('version', 'N/A')}")
|
||||
print(f"find {len(test_cases)} test cases")
|
||||
|
||||
return test_cases
|
||||
except FileNotFoundError:
|
||||
print(f"error: file not found {test_file}")
|
||||
sys.exit(1)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"error: JSON decode failed - {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"-t", "--test_path",
|
||||
required=True,
|
||||
help="Path to test case JSON file"
|
||||
)
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Server host")
|
||||
parser.add_argument("--port", type=int, default=8080, help="Server port")
|
||||
parser.add_argument(
|
||||
"--loop-forever",
|
||||
action="store_true",
|
||||
help="Run test cases repeatedly until interrupted"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--interval",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Seconds to sleep between test cases"
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
def start_server(server_path, model_path, case_path, port):
|
||||
"""启动llama-server进程，返回subprocess.Popen对象"""
|
||||
cmd = ["sudo", server_path, "-m", str(model_path), "--case", str(case_path), "--port", str(port), "-np", str(1), "--verbose", "--no-warmup"]
|
||||
print(f"Start server: {' '.join(cmd)}")
|
||||
process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
time.sleep(1) # 稍等片刻让进程启动
|
||||
return process
|
||||
|
||||
def health_check(HEALTH_URL):
|
||||
"""健康检查，重试直到服务就绪或超时"""
|
||||
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 extract_stream_parts(chunk):
|
||||
choices = chunk.get("choices", [])
|
||||
if not choices:
|
||||
return "", ""
|
||||
|
||||
c0 = choices[0]
|
||||
delta = c0.get("delta", {}) or {}
|
||||
|
||||
think_piece = delta.get("reasoning_content", "") or ""
|
||||
answer_piece = delta.get("content", "") or ""
|
||||
|
||||
# 兼容部分服务端在非 delta 里返回
|
||||
if not think_piece and not answer_piece:
|
||||
msg = c0.get("message", {}) or {}
|
||||
think_piece = msg.get("reasoning_content", "") or ""
|
||||
answer_piece = msg.get("content", "") or c0.get("text", "") or ""
|
||||
|
||||
return think_piece, answer_piece
|
||||
|
||||
def send_request(test_case, API_URL):
|
||||
model = "deepseek"
|
||||
test_id = test_case.get("id", "unknown")
|
||||
prompt = test_case.get("prompt", "")
|
||||
if not prompt:
|
||||
return {"error": f"prompt {test_id} is null, skip"}
|
||||
|
||||
request_body = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"stream": True
|
||||
}
|
||||
|
||||
print("Prompt:")
|
||||
print(prompt)
|
||||
print("Response(stream):")
|
||||
|
||||
think_parts = []
|
||||
answer_parts = []
|
||||
think_opened = False
|
||||
answer_opened = False
|
||||
|
||||
usage = {}
|
||||
timings = {}
|
||||
|
||||
try:
|
||||
with requests.post(API_URL, json=request_body, stream=True, timeout=150) as response:
|
||||
response.raise_for_status()
|
||||
response.encoding = "utf-8"
|
||||
|
||||
for raw_line in response.iter_lines(decode_unicode=False):
|
||||
if not raw_line:
|
||||
continue
|
||||
|
||||
line = raw_line.decode("utf-8", errors="replace")
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
|
||||
payload = line[5:].strip()
|
||||
if not payload:
|
||||
continue
|
||||
if payload == "[DONE]":
|
||||
break
|
||||
|
||||
try:
|
||||
chunk = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
think_piece, answer_piece = extract_stream_parts(chunk)
|
||||
|
||||
if think_piece:
|
||||
if not think_opened:
|
||||
print("[THINK] ", end="", flush=True)
|
||||
think_opened = True
|
||||
print(think_piece, end="", flush=True)
|
||||
think_parts.append(think_piece)
|
||||
|
||||
if answer_piece:
|
||||
if not answer_opened:
|
||||
if think_opened:
|
||||
print("\n[ANSWER] ", end="", flush=True)
|
||||
else:
|
||||
print("[ANSWER] ", end="", flush=True)
|
||||
answer_opened = True
|
||||
print(answer_piece, end="", flush=True)
|
||||
answer_parts.append(answer_piece)
|
||||
|
||||
if "usage" in chunk and isinstance(chunk["usage"], dict):
|
||||
usage = chunk["usage"]
|
||||
if "timings" in chunk and isinstance(chunk["timings"], dict):
|
||||
timings = chunk["timings"]
|
||||
|
||||
print()
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
return {"error": f"test case {test_id} timed out"}
|
||||
except requests.exceptions.RequestException as e:
|
||||
return {"error": f"Request failed: {e}"}
|
||||
|
||||
think_content = "".join(think_parts)
|
||||
content = "".join(answer_parts)
|
||||
|
||||
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)
|
||||
|
||||
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,
|
||||
"think_content": think_content,
|
||||
"content": content,
|
||||
"model": model
|
||||
}
|
||||
|
||||
def print_metrics(round_id, i, total, metrics):
|
||||
print("\n" + "=" * 80)
|
||||
print(f"[Round {round_id} | {i}/{total}]")
|
||||
print(
|
||||
f"prompt_tokens={metrics['prompt_tokens']}, "
|
||||
f"completion_tokens={metrics['completion_tokens']}, "
|
||||
f"decode_tps={metrics['decode_tps']:.2f}, "
|
||||
f"ttft={metrics['ttft']:.2f} ms, "
|
||||
f"prefill_tps={metrics['prefill_tps']:.2f}"
|
||||
)
|
||||
print("=" * 80)
|
||||
|
||||
def log_data(metrics):
|
||||
"""记录数据到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", newline="", encoding="utf-8") as f:
|
||||
writer = csv.writer(f)
|
||||
writer.writerow([
|
||||
timestamp,
|
||||
metrics["model"],
|
||||
metrics["prompt_tokens"],
|
||||
metrics["completion_tokens"],
|
||||
f"{metrics['memory_usage']:.2f}",
|
||||
f"{metrics['decode_tps']:.2f}",
|
||||
f"{metrics['ttft']:.2f}",
|
||||
f"{metrics['prefill_tps']:.2f}",
|
||||
f"{metrics['Power']:.1f}",
|
||||
])
|
||||
|
||||
with open(CONTENT_FILE, "a", newline="", encoding="utf-8") as f:
|
||||
writer = csv.writer(f)
|
||||
writer.writerow([
|
||||
timestamp,
|
||||
metrics["model"],
|
||||
metrics["prompt"],
|
||||
metrics["content"],
|
||||
])
|
||||
|
||||
def main():
|
||||
"""主函数，执行完整流程"""
|
||||
print("=== Starting llama-server test ===")
|
||||
server_process = None
|
||||
try:
|
||||
args = parse_args()
|
||||
|
||||
# 1. 初始化日志文件，加载json测试集
|
||||
init_log_file()
|
||||
|
||||
test_cases = load_test_cases(args.test_path)
|
||||
if not test_cases:
|
||||
print("error:load test_cases failed")
|
||||
sys.exit(1)
|
||||
|
||||
HEALTH_URL = f"http://{args.host}:{args.port}/health"
|
||||
API_URL = f"http://{args.host}:{args.port}/v1/chat/completions"
|
||||
|
||||
# server_process = start_server(args.server_path, MODEL_PATH, CALBIN_PATH, PORT)
|
||||
|
||||
# 3. 健康检查
|
||||
if not health_check(HEALTH_URL):
|
||||
raise RuntimeError("Health check failed")
|
||||
|
||||
round_id = 0
|
||||
|
||||
while True:
|
||||
round_id += 1
|
||||
successful_tests = 0
|
||||
# 每轮发送次数，默认=测试集大小；可按需改成固定值
|
||||
sample_count = len(test_cases)
|
||||
|
||||
print(f"\n=== Round {round_id} start: {sample_count} samples (with replacement) ===")
|
||||
|
||||
for i in range(1, sample_count + 1):
|
||||
test_case = random.choice(test_cases) # 有放回抽样，允许重复
|
||||
print(f"\n[Round {round_id} | Progress: {i}/{sample_count}]")
|
||||
|
||||
metrics = send_request(test_case, API_URL)
|
||||
if not metrics or "error" in metrics:
|
||||
err = metrics["error"] if isinstance(metrics, dict) and "error" in metrics else "unknown error"
|
||||
print(f"[Round {round_id} | {i}/{sample_count}] ERROR: {err}")
|
||||
else:
|
||||
print_metrics(round_id, i, sample_count, metrics)
|
||||
log_data(metrics)
|
||||
successful_tests += 1
|
||||
|
||||
if args.interval > 0 and i < sample_count:
|
||||
time.sleep(args.interval)
|
||||
|
||||
print(f"=== Round {round_id} done: {successful_tests}/{sample_count} success ===")
|
||||
|
||||
if not args.loop_forever:
|
||||
break
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\nInterrupted by user, stopping stress test...")
|
||||
except Exception as e:
|
||||
print(f"✘ Test failed: {e}")
|
||||
sys.exit(1)
|
||||
finally:
|
||||
# 6. 停止服务器进程
|
||||
if server_process and server_process.poll() is None: # 检查进程是否仍在运行
|
||||
server_process.terminate()
|
||||
try:
|
||||
server_process.wait(timeout=5)
|
||||
print("✔ Test completed")
|
||||
except subprocess.TimeoutExpired:
|
||||
server_process.kill()
|
||||
print("✔ Test completed")
|
||||
|
||||
sys.exit(0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user