在 Async Python 中进行 CPU 密集型工作

from your_app.background import task_manager

@task_manager.task(name="data_sync", cpu_bound=False)
async def sync_data_task():
    """I/O密集型数据同步任务"""
    while True:
        try:
            # 异步I/O操作
            await asyncio.sleep(300)  # 每5分钟执行一次
        except asyncio.CancelledError:
            break

@task_manager.task(name="image_processor", cpu_bound=True)
async def process_images_task():
    """CPU密集型图像处理任务"""
    while True:
        try:
            # 获取需要处理的图像
            images = await get_pending_images()
            
            # 对每个图像使用进程池处理
            loop = asyncio.get_running_loop()
            for image in images:
                result = await loop.run_in_executor(
                    task_manager.get_process_pool(),
                    process_image_data,  # CPU密集型函数
                    image.data
                )
                await save_result(image.id, result)
                
            await asyncio.sleep(60)
        except asyncio.CancelledError:
            break

import os
import asyncio
import inspect
import logging
from typing import List, Dict, Any, Optional
from multiprocessing import cpu_count
from concurrent.futures import ProcessPoolExecutor
from fastapi import FastAPI

class BackgroundTaskManager:
    """背景任务管理器 - 使用装饰器注册和管理后台任务"""
    
    def __init__(self):
        self.tasks: List[Dict[str, Any]] = []
        self.running_tasks: List[asyncio.Task] = []
        self.process_pool: Optional[ProcessPoolExecutor] = None
    
    def task(self, *, name: Optional[str] = None, cpu_bound: bool = False,
             startup: bool = True, shutdown_timeout: float = 10.0):
        """
        装饰器: 注册一个后台任务
        
        参数:
            name: 任务名称 (默认使用函数名)
            cpu_bound: 是否是CPU密集型任务
            startup: 是否在应用启动时自动启动任务
            shutdown_timeout: 关闭时等待任务完成的超时时间(秒)
        """
        def decorator(func):
            if not asyncio.iscoroutinefunction(func):
                raise ValueError(f"Background task {func.__name__} must be an async function")
            
            task_name = name or func.__name__
            
            # 存储任务信息
            self.tasks.append({
                "name": task_name,
                "function": func,
                "cpu_bound": cpu_bound,
                "startup": startup,
                "shutdown_timeout": shutdown_timeout
            })
            
            # 返回原始函数不变
            return func
            
        return decorator
    
    def setup_process_pool(self):
        """设置进程池 - 用于CPU密集型任务"""
        if any(task["cpu_bound"] for task in self.tasks):
            cpus = get_cpu_limit()
            workers = max(1, cpus - 1)  # 保留一个CPU给事件循环
            self.process_pool = ProcessPoolExecutor(max_workers=workers)
            logging.info(f"Process pool created with {workers} workers (detected {cpus} CPUs)")
    
    async def start_tasks(self, app: FastAPI):
        """启动所有注册的任务"""
        for task_info in self.tasks:
            if task_info["startup"]:
                name = task_info["name"]
                func = task_info["function"]
                
                # 检查函数签名，如果函数接受app参数，则传入
                sig = inspect.signature(func)
                if "app" in sig.parameters:
                    task_coroutine = func(app)
                else:
                    task_coroutine = func()
                
                # 创建并保存任务
                asyncio_task = asyncio.create_task(task_coroutine, name=name)
                self.running_tasks.append(asyncio_task)
                logging.info(f"Started background task: {name}")
    
    async def stop_tasks(self):
        """停止所有运行中的任务"""
        if not self.running_tasks:
            return
            
        # 取消所有任务
        for task in self.running_tasks:
            if not task.done():
                task.cancel()
        
        # 等待所有任务完成或超时
        pending = self.running_tasks
        self.running_tasks = []
        
        if pending:
            done, pending = await asyncio.wait(
                pending, 
                timeout=10.0,
                return_when=asyncio.ALL_COMPLETED
            )
            
            # 记录未完成的任务
            if pending:
                task_names = [task.get_name() for task in pending]
                logging.warning(f"Some background tasks did not complete in time: {task_names}")
    
    def shutdown_process_pool(self):
        """关闭进程池"""
        if self.process_pool:
            self.process_pool.shutdown(wait=True)
            self.process_pool = None
            logging.info("Process pool shut down")
    
    def get_process_pool(self):
        """获取进程池 - 用于CPU密集型任务"""
        return self.process_pool

def get_cpu_limit():
    """获取容器环境或物理机中可用的CPU数量"""
    try:
        # 检查cgroups v1
        if os.path.exists("/sys/fs/cgroup/cpu/cpu.cfs_quota_us"):
            with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as fp:
                cfs_quota_us = int(fp.read().strip())
            if cfs_quota_us > 0:
                with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as fp:
                    cfs_period_us = int(fp.read().strip())
                container_cpus = cfs_quota_us // cfs_period_us
                return container_cpus
        
        # 检查cgroups v2
        elif os.path.exists("/sys/fs/cgroup/cpu.max"):
            with open("/sys/fs/cgroup/cpu.max") as fp:
                content = fp.read().strip().split()
                if content[0] != "max":
                    quota = int(content[0])
                    period = int(content[1])
                    container_cpus = quota // period
                    return container_cpus
    
    except Exception as e:
        logging.warning(f"Error detecting container CPU limit: {e}")
    
    # 回退到系统CPU计数
    return cpu_count()

@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    # 设置进程池
    task_manager.setup_process_pool()
    
    # 将进程池添加到应用状态
    app.state.process_pool = task_manager.get_process_pool()
    
    # 启动所有注册的后台任务
    await task_manager.start_tasks(app)
    
    # 其他初始化...
    
    yield
    
    # 停止所有后台任务
    await task_manager.stop_tasks()
    
    # 关闭进程池
    task_manager.shutdown_process_pool()
    
    # 其他清理...

app = FastAPI(lifespan=lifespan)