Python多进程、协程混合编程

import asyncio
import multiprocessing as mp
import time
from concurrent.futures.process import ProcessPoolExecutor
from typing import Dict, Any, List

import aiohttp

# ====================== 1. 全局配置管控 ======================
# 切换此开关测试不同的海量数据处理模式
USE_STREAMING = True  # True：开启无锁异步队列流式控制（适合 10万+ 任务）；False：批量开发

AIOHTTP_CONCURRENCY_LIMIT = 10  # 协程高并发连接限制
REQ_TIMEOUT = aiohttp.ClientTimeout(total=5)  # 网络请求超时设置

# Python 跑 CPU 密集型多进程坚决不超卖物理核，留 1 个核给主协程做高并发网络调度
CPU_WORKERS = max(1, mp.cpu_count() - 1)
CPU_PARSE_TIMEOUT = 10.0  # 单个 CPU 任务超时熔断


# ====================== 2. CPU 密集同步计算 ======================
# 注意：在 spawn 启动模式下，被进程池调用的函数必须放在模块顶层
def cpu_intensive_parse(raw_data: Dict[str, Any]) -> str:
    total = 0
    # 模拟 OpenCV 切图、文本区域清洗或密集 JSON 组装
    for i in range(1_000_000):
        total += i * i
    return f"Parsed_Result_{raw_data['url']}_{total}"


# ====================== 3. 单个任务的混合流水线 ======================
async def fetch_and_parse(
        session: aiohttp.ClientSession,
        cpu_pool: ProcessPoolExecutor,
        url: str
) -> Dict[str, Any]:
    try:
        # ---------------- 阶段 A: I/O 密集型操作 ----------------
        # 此时当前协程挂起，Event Loop 去处理其他协程，完全非阻塞
        async with session.get(url, timeout=REQ_TIMEOUT) as resp:
            await resp.text()
            raw_data = {'url': url, "status": resp.status, "data": "raw_html..."}

        if raw_data["status"] != 200:
            return raw_data

        # ---------------- 阶段 B: CPU 密集型操作 ----------------
        # 1. 同步提交到全局进程池，立即拿到 concurrent.futures.Future
        pool_future = cpu_pool.submit(cpu_intensive_parse, raw_data)

        # 2. 核心：使用 asyncio.wrap_future 将进程池的 Future 桥接为协程可以 await 的对象
        # 3. 配合 asyncio.wait_for 实现纯异步的超时熔断，在此期间主事件循环处于完全活跃状态
        parsed_res = await asyncio.wait_for(
            asyncio.wrap_future(pool_future),
            timeout=CPU_PARSE_TIMEOUT
        )
        return {"url": url, "status": 200, "parsed": parsed_res}
    except asyncio.TimeoutError:
        # 专门捕获阶段 B 的 CPU 计算超时
        return {"url": url, "status": "failed", "error": f"CPU 解析超时超过 {CPU_PARSE_TIMEOUT}s"}
    except Exception as e:
        # 捕获网络、DNS、各种突发异常，确保单条任务失败不影响整条流水线
        return {"url": url, "status": "failed", "error": str(e)}


# ====================== 4. 方案 A：流式消费者工人 ======================
async def queue_consumer_worker(
        queue: asyncio.Queue,
        session: aiohttp.ClientSession,
        cpu_pool: ProcessPoolExecutor,
        final_results: List[Dict[str, Any]]
):
    """
    流式工人协程：只要队列里还有 URL 就源源不断地拿出来消费
    并发的工人数量完全由 AIOHTTP_CONCURRENCY_LIMIT 控制，内存始终恒定
    """
    while True:
        try:
            # 非阻塞地从队列获取一个 URL
            url = queue.get_nowait()
        except asyncio.QueueEmpty:
            # 队列空了，工人优雅退出
            break

        # 执行完整的 I/O + CPU 流水线
        result = await fetch_and_parse(session, cpu_pool, url)
        final_results.append(result)

        # 告诉队列该任务处理完毕
        queue.task_done()


# ====================== 5. 异步主编排入口 ======================
async def main_async():
    # 模拟 20 个测试 url
    urls = [f"https://www.baidu.com" for i in range(20)]

    # 初始化全局唯一进程池
    cpu_pool = ProcessPoolExecutor(max_workers=CPU_WORKERS)

    # 统一管控 HTTP 连接池
    connector = aiohttp.TCPConnector(limit=AIOHTTP_CONCURRENCY_LIMIT)

    start_time = time.time()
    final_results = []

    print(f"[主协程] 启动 | 物理核心数: {mp.cpu_count()} | 分配进程池大小: {CPU_WORKERS}")
    print(f"[主协程] 当前运行模式: {'方案A: 异步队列流式' if USE_STREAMING else '方案B: 批量并发'}")

    async with aiohttp.ClientSession(connector=connector) as session:
        if USE_STREAMING:
            # ---------------- 方案 A：纯正的异步队列流式处理 ----------------
            # 1. 初始化异步队列，并将文本装载进队列（只存字符串，10万条也只占几MB内存）
            task_queue = asyncio.Queue()
            for url in urls:
                task_queue.put_nowait(url)

            # 2. 同时启动固定数量的消费者工人
            workers = [
                asyncio.create_task(queue_consumer_worker(task_queue, session, cpu_pool, final_results))
                for _ in range(AIOHTTP_CONCURRENCY_LIMIT)
            ]

            # 3. 挂起等待所有工人把队列里的任务全部啃完
            await asyncio.gather(*workers)
        else:
            # ---------------- 方案 B：简单优雅的批量并发 ----------------
            tasks = [fetch_and_parse(session, cpu_pool, url) for url in urls]
            # gather 会并发调度所有任务，并严格保持原数组顺序返回
            final_results = await asyncio.gather(*tasks, return_exceptions=True)

    # 优雅关闭进程池，cancel_futures=True 会自动取消队列中还没有来得及运行的物理进程任务
    cpu_pool.shutdown(wait=True, cancel_futures=True)

    # 结果统计展示
    success_count = sum(1 for r in final_results if isinstance(r, dict) and r.get("status") == 200)

    print("=" * 60)
    print(f"[主协程] 全部任务执行完成！")
    print(f"总处理条目: {len(final_results)} | 成功解析数: {success_count}")
    print(f"系统总耗时: {time.time() - start_time:.2f} 秒")
    return final_results


# ====================== 6. 顶层物理进程启动保护 ======================
if __name__ == "__main__":
    # 强制使用 spawn 启动模式。在对接 vLLM、PyTorch 或 CUDA 时，
    # 这是唯一能防止 CUDA 运行时死锁、守护进程分叉崩溃的安全模式。
    mp.set_start_method("spawn", force=True)

    # 启动 Asyncio 顶层事件循环
    asyncio.run(main_async())