101 lines
3.9 KiB
Python
101 lines
3.9 KiB
Python
"""并发调用多个下游（例如大模型 / 检索 / 第三方 API）：并发上限、单次超时、有限重试、部分失败、整体截止时间。
|
|
|
|
只用标准库。运行：python fanout.py；测试：python -m pytest -q test_fanout.py
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import random
|
|
import time
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass
|
|
from typing import Generic, TypeVar
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class RetryableError(Exception):
|
|
"""可以重试的错误（限流 429、服务暂时不可用 503 等）。"""
|
|
|
|
|
|
@dataclass
|
|
class Result(Generic[T]):
|
|
key: str
|
|
value: T | None = None
|
|
error: str | None = None
|
|
attempts: int = 0
|
|
|
|
|
|
async def call_with_retry(fn: Callable[[], Awaitable[T]], *, timeout: float, retries: int,
|
|
base_delay: float = 0.05) -> tuple[T, int]:
|
|
"""单次调用带超时；遇到超时或可重试错误时指数退避重试。返回 (结果, 尝试次数)。"""
|
|
for attempt in range(1, retries + 2):
|
|
try:
|
|
async with asyncio.timeout(timeout): # Python 3.11+
|
|
return await fn(), attempt
|
|
except (TimeoutError, RetryableError):
|
|
if attempt > retries:
|
|
raise
|
|
await asyncio.sleep(base_delay * 2 ** (attempt - 1) * (1 + random.random() / 2))
|
|
raise AssertionError("不会执行到这里")
|
|
|
|
|
|
async def fan_out(jobs: dict[str, Callable[[], Awaitable[T]]], *, concurrency: int, timeout: float,
|
|
retries: int = 1, deadline: float | None = None) -> list[Result[T]]:
|
|
"""并发执行所有任务，最多 concurrency 个同时进行；单个失败不影响其他任务。
|
|
deadline：整体截止时间（秒），到时未完成的任务被取消并标记为超时。"""
|
|
sem = asyncio.Semaphore(concurrency)
|
|
results = {k: Result[T](key=k) for k in jobs}
|
|
|
|
async def run(key: str, fn: Callable[[], Awaitable[T]]) -> None:
|
|
async def counted() -> T: # 记录尝试次数（失败时也能看到试了几次）
|
|
results[key].attempts += 1
|
|
return await fn()
|
|
|
|
async with sem:
|
|
try:
|
|
results[key].value, _ = await call_with_retry(counted, timeout=timeout, retries=retries)
|
|
except asyncio.CancelledError:
|
|
results[key].error = "整体超时，已取消"
|
|
raise # 取消必须继续向上传播
|
|
except Exception as e: # 单个任务失败只记录，不影响其他任务
|
|
results[key].error = f"{type(e).__name__}: {e}" if str(e) else type(e).__name__
|
|
|
|
tasks = [asyncio.create_task(run(k, fn)) for k, fn in jobs.items()]
|
|
done, pending = await asyncio.wait(tasks, timeout=deadline)
|
|
for t in pending:
|
|
t.cancel()
|
|
await asyncio.gather(*pending, return_exceptions=True) # 等待取消完成，避免任务泄漏
|
|
return [results[k] for k in jobs]
|
|
|
|
|
|
# ---------------- 演示 ----------------
|
|
def fake_llm(name: str, delay: float, fail_times: int = 0) -> Callable[[], Awaitable[str]]:
|
|
state = {"left": fail_times}
|
|
|
|
async def call() -> str:
|
|
await asyncio.sleep(delay)
|
|
if state["left"] > 0:
|
|
state["left"] -= 1
|
|
raise RetryableError("429 Too Many Requests")
|
|
return f"{name} 的回答"
|
|
|
|
return call
|
|
|
|
|
|
async def demo() -> None:
|
|
jobs = {
|
|
"摘要": fake_llm("摘要", 0.10),
|
|
"分类": fake_llm("分类", 0.10, fail_times=1), # 第一次被限流，重试后成功
|
|
"翻译": fake_llm("翻译", 2.00), # 太慢，单次超时
|
|
"关键词": fake_llm("关键词", 0.05),
|
|
}
|
|
start = time.perf_counter()
|
|
for r in await fan_out(jobs, concurrency=2, timeout=0.5, retries=1):
|
|
print(f"{r.key}: {r.value or r.error}（尝试 {r.attempts} 次）")
|
|
print(f"总耗时 {time.perf_counter() - start:.2f}s")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(demo())
|