iomgaa
157a27f3bb
The telemetry write budget needs asyncio.timeout, whose uncancel accounting was only fixed after 3.11.1 — pinning the floor at 3.12 removes that hazard instead of working around it. Raising ruff's target-version turns on UP047, so gather_bounded, _anext_within and stream_with_liveness_timeouts move to def f[T](...) and the two module-level TypeVars go away. That syntax is a SyntaxError on 3.11, so it can only land together with the version bump.
106 lines
3.9 KiB
Python
106 lines
3.9 KiB
Python
"""三层流式活性超时看门狗: 给任意异步迭代器施加首产出/产出间/总时长三层超时保护。
|
|
|
|
领域无关、不打日志、不触网: 仅以 asyncio.timeout 包裹取值,便于注入假异步
|
|
迭代器做单测。超时抛 StreamLivenessTimeout(携 kind/elapsed/是否已出首项),
|
|
由 transport 边界翻译为 TransientError 并埋点。
|
|
|
|
移植自 Video-Tree-TRM5 adapters/streaming.py(三项目同款,语义零改动)。
|
|
活性口径: 看门狗只感知"有无产出",不区分产出内容——transport 把
|
|
reasoning_content 增量同样作为流元素产出,思考流因此天然刷新计时(CHS R1)。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import time
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator
|
|
|
|
|
|
class StreamLivenessTimeout(Exception): # noqa: N818 — 三项目冻结的公共名
|
|
"""流活性超时异常。
|
|
|
|
属性:
|
|
kind: 超时类型,取值 {"ttft", "inter_token", "total"}。
|
|
elapsed_s: 自流开始到触发超时的秒数。
|
|
first_token_seen: 触发时是否已产出过首项。
|
|
"""
|
|
|
|
def __init__(self, kind: str, elapsed_s: float, first_token_seen: bool) -> None:
|
|
self.kind = kind
|
|
self.elapsed_s = elapsed_s
|
|
self.first_token_seen = first_token_seen
|
|
super().__init__(f"流活性超时({kind}, elapsed={elapsed_s:.1f}s)")
|
|
|
|
|
|
async def _anext_within[T](
|
|
it: AsyncIterator[T],
|
|
timeout_s: float,
|
|
*,
|
|
kind: str,
|
|
start: float,
|
|
first: bool,
|
|
) -> T:
|
|
"""限时取下一项;本层 deadline 触发抛 StreamLivenessTimeout(kind)。
|
|
|
|
上游自抛的 TimeoutError 用 cm.expired() 区分,原样上抛不误吞。
|
|
"""
|
|
try:
|
|
async with asyncio.timeout(timeout_s) as cm:
|
|
return await it.__anext__()
|
|
except TimeoutError:
|
|
if not cm.expired():
|
|
raise # 上游自抛的 TimeoutError,非本层 deadline,原样上抛
|
|
raise StreamLivenessTimeout(kind, time.monotonic() - start, not first) from None
|
|
|
|
|
|
async def stream_with_liveness_timeouts[T](
|
|
source: AsyncIterator[T],
|
|
*,
|
|
ttft_s: float,
|
|
inter_token_s: float,
|
|
total_s: float,
|
|
) -> AsyncIterator[T]:
|
|
"""逐项产出 source,并施加三层活性超时。
|
|
|
|
关键实现: 超时**只包裹单次 __anext__**,绝不包裹 yield——否则总时长
|
|
deadline 会在生成器挂起(消费者处理已产出项)期间继续计时,可能取消
|
|
消费者的 await 而非由本生成器稳定抛出(取消泄漏)。总时长由「每轮取
|
|
min(本层预算, 剩余总额) 钳制」实现: 钳到剩余总额时若超时即判 total,
|
|
否则判 ttft/inter_token。finally 关闭底层迭代器释放上游资源。
|
|
"""
|
|
it = source.__aiter__()
|
|
start = time.monotonic()
|
|
deadline = start + total_s
|
|
first = True
|
|
try:
|
|
while True:
|
|
remaining_total = deadline - time.monotonic()
|
|
if remaining_total <= 0:
|
|
raise StreamLivenessTimeout("total", time.monotonic() - start, not first)
|
|
budget = ttft_s if first else inter_token_s
|
|
# 剩余总额比本层预算更紧 → 先到的是 total
|
|
clamped = remaining_total <= budget
|
|
kind = "total" if clamped else ("ttft" if first else "inter_token")
|
|
try:
|
|
item = await _anext_within(
|
|
it,
|
|
remaining_total if clamped else budget,
|
|
kind=kind,
|
|
start=start,
|
|
first=first,
|
|
)
|
|
except StopAsyncIteration:
|
|
return
|
|
first = False
|
|
yield item
|
|
finally:
|
|
# 收尾关闭上游迭代器;清理失败不得掩盖主异常,故显式抑制
|
|
aclose = getattr(it, "aclose", None)
|
|
if aclose is not None:
|
|
with contextlib.suppress(Exception):
|
|
await aclose()
|