import asyncio from typing import Callable, Coroutine, List, Any, Dict async def run_tasks_with_async_worker_group( task_list: List[Any], handler: Callable[[Any], Coroutine[Any, Any, None]], *, description: str = None, unit: str = "item", max_concurrency: int = 20, fail_fast: bool = False, show_progress: bool = True, ) -> Dict[str, Any]: """使用异步 worker pool 处理 I/O 密集型任务 Args: task_list: 任务对象列表 handler: 异步处理函数,接受单个任务对象 description: 进度条描述 unit: 进度条单位 max_concurrency: 最大并发 worker 数 fail_fast: True 时任一 worker 异常立即中止全部 show_progress: 是否显示 tqdm 进度条(需要安装 tqdm) Returns: {"total_task": int, "processed_task": int, "errors": [(idx, obj, exc), ...]} """ if not task_list: return {"total_task": 0, "processed_task": 0, "errors": []} total_task = len(task_list) processed_task = 0 errors: List[tuple] = [] queue: asyncio.Queue = asyncio.Queue() cancel_event = asyncio.Event() counter_lock = asyncio.Lock() # 进度条(可选) if show_progress: try: from tqdm.asyncio import tqdm as async_tqdm pbar = async_tqdm(total=total_task, unit=unit, desc=description) except ImportError: pbar = None else: pbar = None async def worker(worker_id: int): nonlocal processed_task while True: item = await queue.get() if item is None: queue.task_done() break idx, task_obj = item try: if cancel_event.is_set(): return await handler(task_obj) async with counter_lock: processed_task += 1 except Exception as e: if fail_fast: cancel_event.set() raise errors.append((idx, task_obj, e)) finally: if pbar: pbar.update() queue.task_done() workers = [asyncio.create_task(worker(i)) for i in range(max_concurrency)] try: for index, task in enumerate(task_list, start=1): await queue.put((index, task)) await queue.join() except Exception: for w in workers: w.cancel() raise finally: for _ in range(max_concurrency): await queue.put(None) await asyncio.gather(*workers, return_exceptions=True) if pbar: pbar.close() return { "total_task": total_task, "processed_task": processed_task, "errors": errors, } async def run_tasks_with_asyncio_task_group( task_list: List[Any], handler: Callable[[Any], Coroutine[Any, Any, None]], *, description: str = None, unit: str = "item", max_concurrency: int = 20, fail_fast: bool = False, show_progress: bool = True, ) -> Dict[str, Any]: """使用 asyncio.TaskGroup 处理 I/O 密集型任务 Args: task_list: 任务对象列表 handler: 异步处理函数,接受单个任务对象 description: 进度条描述 unit: 进度条单位 max_concurrency: 最大并发数(通过 Semaphore 控制) fail_fast: True 时任一任务异常立即中止全部 show_progress: 是否显示 tqdm 进度条 Returns: {"total_task": int, "processed_task": int, "errors": [(idx, obj, exc), ...]} """ if not task_list: return {"total_task": 0, "processed_task": 0, "errors": []} processed_task = 0 total_task = len(task_list) errors: List[tuple] = [] semaphore = asyncio.Semaphore(max_concurrency) # 进度条(可选) if show_progress: try: from tqdm.asyncio import tqdm as async_tqdm pbar = async_tqdm(total=total_task, unit=unit, desc=description) except ImportError: pbar = None else: pbar = None lock = asyncio.Lock() async def _run_single_task(task_obj: Any, idx: int): nonlocal processed_task async with semaphore: try: await handler(task_obj) async with lock: processed_task += 1 except Exception as e: if fail_fast: raise e errors.append((idx, task_obj, e)) finally: if pbar: pbar.update() async with asyncio.TaskGroup() as task_group: for index, task in enumerate(task_list, start=1): task_group.create_task( _run_single_task(task, index), name=f"processing {description}-{index}", ) if pbar: pbar.close() return { "total_task": total_task, "processed_task": processed_task, "errors": errors, }