four-level-schedule/test_concurrent.py

#!/usr/bin/env python3
"""
并发测试脚本 - 测试 parse-service 的并发处理能力
"""
import requests
import json
import time
import threading
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed

# 配置
PARSE_SERVICE_URL = "http://localhost:8000"
EMBEDDING_API_URL = "http://localhost:8084"

# 测试文件列表（可以重复使用来模拟并发）
TEST_FILES = [
    "/Users/chaizi/Work/Projects/yusys/ai-study/code/four-level-schedule/duomotai/examples/5.docx",
    "/Users/chaizi/Work/Projects/yusys/ai-study/code/four-level-schedule/duomotai/examples/20以内口算.pdf",
]


def check_services():
    """检查服务状态"""
    print("=" * 70)
    print("检查服务状态...")
    print("=" * 70)

    # 检查 parse-service
    try:
        response = requests.get(f"{PARSE_SERVICE_URL}/health", timeout=5)
        if response.status_code == 200:
            print("✅ parse-service 正常")
        else:
            print(f"❌ parse-service 异常: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ parse-service 不可达: {e}")
        print("   请先启动: cd parse-service && python main.py")
        return False

    # 检查 schedule-embedding-api
    try:
        response = requests.get(f"{EMBEDDING_API_URL}/actuator/health", timeout=5)
        if response.status_code == 200:
            data = response.json()
            if data.get("status") == "UP":
                print("✅ schedule-embedding-api 正常")
            else:
                print(f"❌ schedule-embedding-api 状态异常: {data}")
                return False
        else:
            print(f"❌ schedule-embedding-api 异常: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ schedule-embedding-api 不可达: {e}")
        return False

    print()
    return True


def parse_single_file(file_path, task_id):
    """解析单个文件（单线程任务）"""
    start_time = time.time()

    try:
        # 步骤 1: 调用 parse-service 解析
        url = f"{PARSE_SERVICE_URL}/api/v1/parse/path?file_path={file_path}"
        response = requests.post(url, json={}, timeout=300)

        if response.status_code != 200:
            return {
                "task_id": task_id,
                "file_path": file_path,
                "success": False,
                "error": f"HTTP error: {response.status_code}",
                "parse_time_ms": 0,
                "index_time_ms": 0,
                "total_time_ms": 0
            }

        data = response.json()
        if data.get("code") != 200:
            return {
                "task_id": task_id,
                "file_path": file_path,
                "success": False,
                "error": data.get("message", "Unknown error"),
                "parse_time_ms": 0,
                "index_time_ms": 0,
                "total_time_ms": 0
            }

        parse_data = data.get("data", {})
        parse_time_ms = parse_data.get("parse_time_ms", 0)
        parse_time = time.time() - start_time

        # 步骤 2: 调用 embedding-api 向量化
        index_start_time = time.time()

        index_url = f"{EMBEDDING_API_URL}/api/v1/documents/index"

        file_name = file_path.split("/")[-1] if "/" in file_path else file_path

        index_request = {
            "docId": task_id,
            "fileName": file_name,
            "fullText": parse_data.get("content", ""),
            "filePath": file_path,
            "fileType": parse_data.get("file_type"),
            "metadata": parse_data.get("metadata", {})
        }

        # 添加 fileSize
        metadata = parse_data.get("metadata", {})
        if metadata and "file_size" in metadata:
            index_request["fileSize"] = metadata["file_size"]

        index_response = requests.post(
            index_url,
            json=index_request,
            headers={"Content-Type": "application/json"},
            timeout=120
        )

        index_time = time.time() - index_start_time
        total_time = time.time() - start_time

        if index_response.status_code != 200:
            return {
                "task_id": task_id,
                "file_path": file_path,
                "success": False,
                "error": f"Index HTTP error: {index_response.status_code}",
                "parse_time_ms": parse_time_ms,
                "index_time_ms": int(index_time * 1000),
                "total_time_ms": int(total_time * 1000)
            }

        index_data = index_response.json()
        index_success = index_data.get("success", False)

        return {
            "task_id": task_id,
            "file_path": file_path,
            "success": index_success,
            "error": None if index_success else index_data.get("message", "Index failed"),
            "parse_time_ms": parse_time_ms,
            "index_time_ms": int(index_time * 1000),
            "total_time_ms": int(total_time * 1000)
        }

    except Exception as e:
        total_time = time.time() - start_time
        return {
            "task_id": task_id,
            "file_path": file_path,
            "success": False,
            "error": str(e),
            "parse_time_ms": 0,
            "index_time_ms": 0,
            "total_time_ms": int(total_time * 1000)
        }


def run_concurrent_test(num_tasks=10, num_workers=5):
    """运行并发测试"""
    print("=" * 70)
    print(f"并发测试: {num_tasks} 个任务, {num_workers} 个工作线程")
    print("=" * 70)

    # 准备任务列表（循环使用测试文件）
    tasks = []
    for i in range(num_tasks):
        file_path = TEST_FILES[i % len(TEST_FILES)]
        task_id = f"concurrent-test-{uuid.uuid4().hex[:8]}"
        tasks.append((file_path, task_id))

    print(f"\n准备执行 {len(tasks)} 个任务...")
    for i, (file_path, task_id) in enumerate(tasks[:5]):  # 只显示前5个
        file_name = file_path.split("/")[-1]
        print(f"  [{task_id}] {file_name}")
    if len(tasks) > 5:
        print(f"  ... 还有 {len(tasks) - 5} 个任务")

    print("\n开始并发执行...")
    print("-" * 70)

    start_time = time.time()
    results = []

    # 使用线程池执行
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        # 提交所有任务
        future_to_task = {
            executor.submit(parse_single_file, file_path, task_id): (file_path, task_id)
            for file_path, task_id in tasks
        }

        # 收集结果
        completed = 0
        for future in as_completed(future_to_task):
            file_path, task_id = future_to_task[future]
            try:
                result = future.result()
                results.append(result)
                completed += 1

                # 显示进度
                file_name = file_path.split("/")[-1]
                status = "✅" if result["success"] else "❌"
                print(f"[{completed}/{num_tasks}] {status} {task_id} - {file_name} "
                      f"(总耗时: {result['total_time_ms']}ms)")

            except Exception as e:
                print(f"❌ 任务异常: {task_id} - {e}")
                results.append({
                    "task_id": task_id,
                    "file_path": file_path,
                    "success": False,
                    "error": str(e),
                    "parse_time_ms": 0,
                    "index_time_ms": 0,
                    "total_time_ms": 0
                })

    total_time = time.time() - start_time

    # 统计结果
    success_count = sum(1 for r in results if r["success"])
    fail_count = len(results) - success_count

    # 计算平均耗时（仅统计成功的任务）
    successful_results = [r for r in results if r["success"]]
    if successful_results:
        avg_parse_time = sum(r["parse_time_ms"] for r in successful_results) / len(successful_results)
        avg_index_time = sum(r["index_time_ms"] for r in successful_results) / len(successful_results)
        avg_total_time = sum(r["total_time_ms"] for r in successful_results) / len(successful_results)
    else:
        avg_parse_time = avg_index_time = avg_total_time = 0

    # 打印总结
    print("\n" + "=" * 70)
    print("并发测试总结")
    print("=" * 70)
    print(f"总任务数: {len(results)}")
    print(f"并发数: {num_workers}")
    print(f"成功: {success_count}")
    print(f"失败: {fail_count}")
    print(f"总耗时: {total_time:.2f} 秒")
    print(f"吞吐量: {len(results)/total_time:.2f} 任务/秒")
    print()
    if successful_results:
        print("平均耗时（仅成功任务）:")
        print(f"  - 解析: {avg_parse_time:.0f} ms")
        print(f"  - 向量化: {avg_index_time:.0f} ms")
        print(f"  - 总计: {avg_total_time:.0f} ms")

    if fail_count > 0:
        print("\n失败任务详情:")
        for r in results:
            if not r["success"]:
                file_name = r["file_path"].split("/")[-1]
                print(f"  - [{r['task_id']}] {file_name}: {r['error']}")

    print("=" * 70)

    return results


def main():
    print("\n")
    print("╔" + "=" * 68 + "╗")
    print("║" + " " * 18 + "多模态解析 - 并发测试" + " " * 24 + "║")
    print("╚" + "=" * 68 + "╝")
    print()

    # 检查服务
    if not check_services():
        return 1

    # 测试配置
    print("请选择测试模式:")
    print("  1. 小规模测试 (5 任务, 2 并发)")
    print("  2. 中规模测试 (10 任务, 5 并发)")
    print("  3. 大规模测试 (20 任务, 10 并发)")
    print("  4. 自定义...")

    try:
        choice = input("\n请输入选项 (1-4, 默认 2): ").strip()

        if choice == "1":
            num_tasks, num_workers = 5, 2
        elif choice == "3":
            num_tasks, num_workers = 20, 10
        elif choice == "4":
            num_tasks = int(input("请输入任务数量: "))
            num_workers = int(input("请输入并发数: "))
        else:
            num_tasks, num_workers = 10, 5

    except KeyboardInterrupt:
        print("\n\n用户中断")
        return 1
    except Exception as e:
        print(f"\n输入错误: {e}，使用默认配置 (10 任务, 5 并发)")
        num_tasks, num_workers = 10, 5

    # 运行测试
    try:
        results = run_concurrent_test(num_tasks, num_workers)
        return 0 if all(r["success"] for r in results) else 1
    except KeyboardInterrupt:
        print("\n\n用户中断")
        return 1


if __name__ == "__main__":
    exit(main())