管理FastAPI中ProcessPoolExecutor的正确姿势

import asyncio
import concurrent.futures
import functools
import re
from fastapi import FastAPI

app = FastAPI()

# 示例：将内容分割成块
def split_on_whitespace(content: str, count: int = 6):
    if not content:
        return ['' for _ in range(count)]
    # 简化示例，实际逻辑可能更复杂
    chunk_size = len(content) // count
    return [content[i:i + chunk_size] for i in range(0, len(content), chunk_size)][:count]

# 示例：在内容块上运行正则表达式
def run_regex_on_content_chunk(content: str):
    domains = []
    domain_patt = re.compile(r'([a-zA-Z0-9\-_]+\.){1,}[a-zA-Z0-9\-_]+')
    for match in domain_patt.finditer(content):
        domains.append(match.group(0))
    return domains

# 辅助函数：在执行器中运行任务
async def executor_task(fn, executor):
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(executor, fn)

@app.post("/addContent")
async def add_content(content: dict):
    all_content = content['data']
    nworkers = 6 # 为每个请求创建6个进程
    content_chunks = split_on_whitespace(all_content, nworkers)
    async_tasks = []

    # 错误：在每个请求中创建新的ProcessPoolExecutor
    with concurrent.futures.ProcessPoolExecutor(max_workers=nworkers) as executor:
        for chunk in content_chunks:
            regex_fn = functools.partial(run_regex_on_content_chunk, chunk)
            async_tasks.append(executor_task(regex_fn, executor))
        results = await asyncio.gather(*async_tasks)

    # 进一步处理results...
    return {"message": "Content processed", "results": results}

# if __name__ == "__main__":
#     import uvicorn
#     uvicorn.run(app, host="0.0.0.0", port=8000)

from contextlib import asynccontextmanager
import asyncio
import concurrent.futures
import functools
import re
from fastapi import FastAPI

# 1. 定义全局进程池变量
process_pool: concurrent.futures.ProcessPoolExecutor | None = None

# 2. 定义一个异步上下文管理器来管理进程池的生命周期
@asynccontextmanager
async def lifespan_event_handler(app: FastAPI):
    global process_pool
    # 建议的工人数量：通常是CPU核心数的1到3倍，具体取决于任务类型和系统负载
    # 对于纯CPU密集型任务，通常不超过CPU核心数。
    # 对于混合型或有少量I/O等待的任务，可以适当增加。
    nworkers = 18 
    process_pool = concurrent.futures.ProcessPoolExecutor(max_workers=nworkers)
    print(f"ProcessPoolExecutor initialized with {nworkers} workers.")
    try:
        yield # FastAPI 将在此处启动服务器并运行应用程序
    finally:
        # 在应用程序关闭时，优雅地关闭进程池
        if process_pool:
            process_pool.shutdown(wait=True)
            print("ProcessPoolExecutor shut down.")

# 3. 在FastAPI应用初始化时，传入lifespan事件处理器
app = FastAPI(lifespan=lifespan_event_handler)

# 示例：将内容分割成块 (与原问题代码相同)
def split_on_whitespace(content: str, count: int = 6):
    if not content:
        return ['' for _ in range(count)]
    chunk_size = len(content) // count
    return [content[i:i + chunk_size] for i in range(0, len(content), chunk_size)][:count]

# 示例：在内容块上运行正则表达式 (与原问题代码相同)
def run_regex_on_content_chunk(content: str):
    domains = []
    domain_patt = re.compile(r'([a-zA-Z0-9\-_]+\.){1,}[a-zA-Z0-9\-_]+')
    for match in domain_patt.finditer(content):
        domains.append(match.group(0))
    return domains

# 辅助函数：在执行器中运行任务 (与原问题代码相同)
async def executor_task(fn, executor):
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(executor, fn)

# 4. 修改API端点，使用全局进程池
@app.post("/addContent")
async def add_content(content: dict):
    global process_pool
    if process_pool is None:
        # 理论上不会发生，因为lifespan确保了它的存在
        raise RuntimeError("ProcessPoolExecutor is not initialized.")

    all_content = content['data']
    # 这里的 nworkers 已经由进程池的 max_workers 决定，
    # 但我们可以根据需要决定分割的块数。
    # 为了简化，这里假设分割成与进程池工人数量相同的块数。
    num_chunks = process_pool._max_workers # 或根据业务逻辑自定义
    content_chunks = split_on_whitespace(all_content, num_chunks)
    async_tasks = []

    for chunk in content_chunks:
        regex_fn = functools.partial(run_regex_on_content_chunk, chunk)
        # 使用全局的 process_pool
        async_tasks.append(executor_task(regex_fn, process_pool))

    results = await asyncio.gather(*async_tasks)

    return {"message": "Content processed successfully", "extracted_domains": results}

# 5. 确保FastAPI应用在主进程中运行
if __name__ == "__main__":
    import uvicorn
    # 启动应用时，lifespan_event_handler 会被调用
    uvicorn.run(app, host="0.0.0.0", port=8000)