feat: add retry middleware with per-attempt governance orchestration
This commit is contained in:
@@ -0,0 +1,295 @@
|
||||
"""RetryMW: 重试循环拥有"尝试"的全部编排(D13 自研;蓝本 CHS governance.py:107-268)。
|
||||
|
||||
每次尝试 = 选源(跳过冷却源)→ 该源熔断门 → 限流 permit → transport;
|
||||
换源、退避、逐次遥测、熔断写回、permit 结算全部在循环内。作为洋葱的
|
||||
终端(compose 的 terminal)被调用: `await retry_mw(request)`。
|
||||
|
||||
例外说明(P7): 本模块 import httpx 仅用于失败原因归类(CHS 同款判定),
|
||||
httpx 是库的核心依赖而非实现层内部件,不违反"middleware 只依赖端口"。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
from loguru import logger
|
||||
|
||||
from polygateway.errors import (
|
||||
AllSourcesExhausted,
|
||||
CircuitOpenError,
|
||||
PolyGatewayError,
|
||||
RequestRejectedError,
|
||||
ResultInvalidError,
|
||||
SourceDeadError,
|
||||
TransientError,
|
||||
)
|
||||
from polygateway.middleware.breaker import BreakerGate
|
||||
from polygateway.middleware.ratelimit import QuotaGate
|
||||
from polygateway.sources import SourceCooldownMemo
|
||||
from polygateway.streaming import StreamLivenessTimeout
|
||||
from polygateway.types import LLMResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from polygateway.ports import (
|
||||
GateDecision,
|
||||
Permit,
|
||||
ProviderGate,
|
||||
RateLimiter,
|
||||
SourceSelector,
|
||||
Transport,
|
||||
)
|
||||
from polygateway.types import (
|
||||
BackpressurePolicy,
|
||||
ChatRequest,
|
||||
RetryPolicy,
|
||||
SourceConfig,
|
||||
TransportResult,
|
||||
)
|
||||
|
||||
|
||||
def _failure_reason(exc: PolyGatewayError) -> str:
|
||||
"""失败原因归类(CHS governance.py:169 同款)。"""
|
||||
if isinstance(exc, SourceDeadError):
|
||||
return "source_dead"
|
||||
if exc.status_code == 429:
|
||||
return "rate_limited"
|
||||
if isinstance(exc.__cause__, httpx.TimeoutException | StreamLivenessTimeout):
|
||||
return "timeout"
|
||||
return "network_error"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Failed:
|
||||
"""一次可重试失败的结果(SourceDead 立即换源,Transient 退避)。"""
|
||||
|
||||
exc: PolyGatewayError
|
||||
immediate: bool
|
||||
|
||||
|
||||
class RetryMW:
|
||||
"""尝试编排器;时钟/睡眠/随机全部注入,纯确定性可测(P6)。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
scope: str,
|
||||
sources: list[SourceConfig],
|
||||
selector: SourceSelector,
|
||||
limiter: RateLimiter,
|
||||
gate: ProviderGate,
|
||||
transport: Transport,
|
||||
retry: RetryPolicy,
|
||||
backpressure: BackpressurePolicy,
|
||||
quota_full: str = "wait",
|
||||
cooldown_memo: SourceCooldownMemo | None = None,
|
||||
emitter: object | None = None,
|
||||
now: Callable[[], float] = time.monotonic,
|
||||
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
rng: Callable[[], float] = random.random,
|
||||
) -> None:
|
||||
if quota_full not in ("wait", "fail_fast"):
|
||||
raise ValueError(f"quota_full 必须是 wait|fail_fast: {quota_full!r}")
|
||||
self._scope = scope
|
||||
self._sources = list(sources)
|
||||
self._selector = selector
|
||||
self._quota = QuotaGate(limiter)
|
||||
self._breaker = BreakerGate(gate)
|
||||
self._transport = transport
|
||||
self._retry = retry
|
||||
self._bp = backpressure
|
||||
self._quota_full = quota_full
|
||||
self._memo = cooldown_memo or SourceCooldownMemo(now=now)
|
||||
self._emitter = emitter
|
||||
self._now = now
|
||||
self._sleep = sleep
|
||||
self._rng = rng
|
||||
|
||||
async def __call__(self, request: ChatRequest) -> LLMResponse:
|
||||
"""执行治理调用;scope 级失败按 §6.1 携结构化字段上抛。"""
|
||||
if not self._sources:
|
||||
raise AllSourcesExhausted(scope=self._scope, reason="no_sources", retry_after_s=0.0)
|
||||
fails = 0
|
||||
reasons: dict[str, str] = {}
|
||||
while True:
|
||||
picked, gate_rejections = await self._pick_runnable(reasons)
|
||||
if picked is None:
|
||||
await self._on_no_runnable(gate_rejections, reasons)
|
||||
continue
|
||||
outcome = await self._attempt(request, *picked, reasons)
|
||||
if isinstance(outcome, LLMResponse):
|
||||
return outcome
|
||||
fails += 1
|
||||
if fails >= self._retry.max_attempts:
|
||||
raise AllSourcesExhausted(
|
||||
scope=self._scope, reason="retry_exhausted",
|
||||
retry_after_s=self._retry.backoff_base_s,
|
||||
per_source_reasons=reasons,
|
||||
) from outcome.exc
|
||||
if not outcome.immediate:
|
||||
await self._sleep(self._backoff_delay(fails, outcome.exc))
|
||||
|
||||
# —— 选源与准入(CHS _pick_runnable 120-167)——
|
||||
|
||||
async def _pick_runnable(
|
||||
self, reasons: dict[str, str]
|
||||
) -> tuple[tuple[SourceConfig, Permit, GateDecision] | None, int]:
|
||||
stats = {s.name: await self._quota.stats(s) for s in self._sources}
|
||||
gate_rejections = 0
|
||||
for cand in self._selector.order(self._sources, stats):
|
||||
if self._memo.active(cand.name):
|
||||
# 冷却备忘跳过也计入拒绝数,保住 circuit_open 判据(CHS 同款)
|
||||
gate_rejections += 1
|
||||
reasons[cand.name] = "cooldown"
|
||||
continue
|
||||
permit = await self._quota.try_acquire(cand)
|
||||
if permit is None:
|
||||
reasons.setdefault(cand.name, "rate_limited")
|
||||
continue
|
||||
try:
|
||||
entry = await self._breaker.try_enter(cand, uuid.uuid4().hex)
|
||||
except BaseException:
|
||||
await self._settle_and_release(permit, 0)
|
||||
raise
|
||||
if entry.allowed:
|
||||
return (cand, permit, entry), gate_rejections
|
||||
gate_rejections += 1
|
||||
reasons[cand.name] = "circuit_open"
|
||||
# 开路源本地记冷却,避免每轮白烧 RPM 探测(CHS governance.py:107)
|
||||
self._memo.set_until(cand.name, self._now() + entry.retry_after_s)
|
||||
await self._settle_and_release(permit, 0)
|
||||
return None, gate_rejections
|
||||
|
||||
async def _on_no_runnable(self, gate_rejections: int, reasons: dict[str, str]) -> None:
|
||||
if gate_rejections == len(self._sources):
|
||||
names = tuple(s.name for s in self._sources)
|
||||
raise CircuitOpenError(
|
||||
scope=self._scope,
|
||||
retry_after_s=await self._breaker.retry_after_s(names),
|
||||
per_source_reasons=reasons,
|
||||
)
|
||||
if self._quota_full == "fail_fast":
|
||||
raise AllSourcesExhausted(
|
||||
scope=self._scope, reason="quota_exhausted",
|
||||
retry_after_s=self._bp.poll_interval_s, per_source_reasons=reasons,
|
||||
)
|
||||
await self._sleep(self._bp.poll_interval_s)
|
||||
|
||||
# —— 单次尝试(CHS run 200-268)——
|
||||
|
||||
async def _attempt(
|
||||
self,
|
||||
request: ChatRequest,
|
||||
source: SourceConfig,
|
||||
permit: Permit,
|
||||
entry: GateDecision,
|
||||
reasons: dict[str, str],
|
||||
) -> LLMResponse | _Failed:
|
||||
call_id = str(uuid.uuid4())
|
||||
started = self._now()
|
||||
actual = 0
|
||||
try:
|
||||
result = await self._transport.complete(
|
||||
messages=request.messages, source=source,
|
||||
stream=request.stream, overlay=request.overlay, call_id=call_id,
|
||||
)
|
||||
actual = result.prompt_tokens + result.completion_tokens
|
||||
await self._breaker.record_success(entry)
|
||||
await self._quota.mark_progress()
|
||||
response = self._build_response(source, result, call_id, started)
|
||||
await self._emit(request, source, call_id, started, response=response)
|
||||
return response
|
||||
except RequestRejectedError as exc:
|
||||
await self._on_rejected(exc, source, entry)
|
||||
await self._emit(request, source, call_id, started, error=exc)
|
||||
raise
|
||||
except ResultInvalidError as exc:
|
||||
# 坏结果 ≠ 坏服务: 熔断记成功,异常上抛消耗业务失败预算(§6.3)
|
||||
await self._breaker.record_success(entry)
|
||||
await self._emit(request, source, call_id, started, error=exc)
|
||||
raise
|
||||
except asyncio.CancelledError:
|
||||
if entry.is_probe:
|
||||
await self._breaker.release_probe(entry)
|
||||
await self._emit(request, source, call_id, started, error="cancelled")
|
||||
raise
|
||||
except (SourceDeadError, TransientError) as exc:
|
||||
dead = isinstance(exc, SourceDeadError)
|
||||
reason = _failure_reason(exc)
|
||||
reasons[source.name] = reason
|
||||
await self._breaker.record_failure(entry, reason, dead)
|
||||
if not dead:
|
||||
actual = source.est_tokens # 保守: 失败请求可能已被网关计费(CHS 同款)
|
||||
await self._emit(request, source, call_id, started, error=exc)
|
||||
return _Failed(exc, immediate=dead)
|
||||
finally:
|
||||
await self._settle_and_release(permit, actual)
|
||||
|
||||
async def _on_rejected(
|
||||
self, exc: RequestRejectedError, source: SourceConfig, entry: GateDecision
|
||||
) -> None:
|
||||
provider_responded = exc.source_name == source.name and exc.status_code is not None
|
||||
if provider_responded:
|
||||
await self._breaker.record_success(entry) # 网关健康地拒了坏请求
|
||||
elif entry.is_probe:
|
||||
await self._breaker.release_probe(entry)
|
||||
|
||||
# —— 辅助 ——
|
||||
|
||||
def _backoff_delay(self, fails: int, exc: PolyGatewayError) -> float:
|
||||
"""指数退避+jitter,与 Retry-After 提示取大(ARCH §7.2;VT jitter 系数)。"""
|
||||
base = min(self._retry.backoff_base_s * (2 ** (fails - 1)), self._retry.backoff_max_s)
|
||||
delay = base * (0.5 + self._rng())
|
||||
retry_after = getattr(exc, "retry_after_s", None) or 0.0
|
||||
return max(delay, retry_after)
|
||||
|
||||
def _build_response(
|
||||
self, source: SourceConfig, result: TransportResult, call_id: str, started: float
|
||||
) -> LLMResponse:
|
||||
return LLMResponse(
|
||||
content=result.content, thinking=result.thinking,
|
||||
model=source.model, provider=source.provider,
|
||||
prompt_tokens=result.prompt_tokens, completion_tokens=result.completion_tokens,
|
||||
latency_ms=int((self._now() - started) * 1000),
|
||||
ttft_ms=result.ttft_ms, max_inter_token_ms=result.max_inter_token_ms,
|
||||
cache_hit=False, call_id=call_id,
|
||||
source_name=source.name, cost=None, usage_source=result.usage_source,
|
||||
)
|
||||
|
||||
async def _settle_and_release(self, permit: Permit, actual: int) -> None:
|
||||
"""finally 专用: settle 后必 release;失败降级 warning,绝不掩盖主异常/取消。"""
|
||||
try:
|
||||
try:
|
||||
await permit.settle(actual)
|
||||
finally:
|
||||
await permit.release()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.warning("permit 结算/释放失败(不掩盖主异常): {}", exc)
|
||||
|
||||
async def _emit(
|
||||
self, request: ChatRequest, source: SourceConfig, call_id: str, started: float,
|
||||
*, response: LLMResponse | None = None, error: object | None = None,
|
||||
) -> None:
|
||||
"""逐次遥测(经注入的单一 Emitter);遥测失败不得影响调用(铁律)。"""
|
||||
if self._emitter is None:
|
||||
return
|
||||
try:
|
||||
await self._emitter.emit_attempt(
|
||||
request=request, source=source, call_id=call_id,
|
||||
latency_ms=int((self._now() - started) * 1000),
|
||||
response=response,
|
||||
error=None if error is None else str(error),
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.warning("逐次遥测记录失败(降级不冒泡): {}", exc)
|
||||
Reference in New Issue
Block a user