← 中文资源

简体中文版本

Python asyncio TaskGroup 批量任务

用 TaskGroup 并发运行协程:任一任务失败即取消同组兄弟任务,并通过 ExceptionGroup 汇总异常。

阅读英文原文

功能说明

asyncio.TaskGroup 是 Python 3.11 引入的结构化并发原语。把多个协程交给同一个任务组后,代码会等待全部任务结束;只要其中一个抛出异常,任务组就会取消其余仍在运行的任务,并把所有异常包装成 ExceptionGroup 抛出。相比手工维护任务列表再调用 gather,这种写法不会遗留无人等待的孤儿任务,异步栈中的错误也更容易定位。

本例把业务任务封装为 Job 数据类,run_batch 在任务组内为每个 Job 创建任务,并把名称与任务句柄一起保存。任务组退出时所有任务已经完成,因此可以直接读取 task.result() 组装结果字典。main 先运行两个正常任务,再运行一个必然失败的批次,用 except* 语法筛选 RuntimeError。

关键代码

run_job 通过 asyncio.sleep 模拟耗时,并在名称为 bad 时抛出 RuntimeError。run_batch 使用 async with asyncio.TaskGroup() as group 创建作用域,循环调用 group.create_task 而不是 asyncio.create_task,这样任务自动归属该组。退出 with 块时隐式等待,任何子任务异常都会在退出点重新抛出。tasks 列表保存任务句柄,退出后可安全读取结果。main 中的 except* 只捕获异常组内的 RuntimeError,其余类型仍会继续传播,便于分层处理。

使用注意

不要吞掉异常组:记录日志后应决定是重试整批还是只重试失败项,否则故障会被静默掩盖。任务组取消是协作式的,子协程必须在可取消的 await 点上运行;长时间占用 CPU 的同步函数不会响应取消,应放入线程池。每个任务都应设置自身的超时,否则一个永久挂起的任务会让整个任务组无法退出。结果字典依赖任务完成顺序之外的名称映射,避免用位置索引对应任务。Python 3.10 及更早版本没有 TaskGroup,需要改用 gather 并自行处理取消。

验证方式

直接运行脚本,第一批输出 alpha 与 beta 两个 ok,说明并发任务结果被正确收集。第二批会在约十毫秒后失败,输出 failures 列表包含 bad failed,随后打印 slow job cancelled by TaskGroup,证明同组任务在首次失败后被取消,而不是等待一秒。可再添加一个正常任务验证它同样被取消。

代码示例

import asyncio
from dataclasses import dataclass


@dataclass
class Job:
    name: str
    seconds: float


async def run_job(job: Job) -> tuple[str, str]:
    await asyncio.sleep(job.seconds)
    if job.name == "bad":
        raise RuntimeError(f"job {job.name} failed")
    return job.name, "ok"


async def run_batch(jobs: list[Job]) -> dict[str, str]:
    tasks: list[tuple[str, asyncio.Task[tuple[str, str]]]] = []
    async with asyncio.TaskGroup() as group:
        for job in jobs:
            tasks.append((job.name, group.create_task(run_job(job))))
    return {name: task.result()[1] for name, task in tasks}


async def main() -> None:
    print(await run_batch([Job("alpha", 0.05), Job("beta", 0.01)]))
    try:
        await run_batch([Job("slow", 1.0), Job("bad", 0.01)])
    except* RuntimeError as group:
        print("failures:", [str(error) for error in group.exceptions])
        print("slow job cancelled by TaskGroup")


asyncio.run(main())