refactor: share one admission path across the three governance loops

_pick_runnable and _on_no_runnable lived in three copies (retry.py,
embedding.py, ocr.py), the latter two being verbatim subsets of the
first. Admission semantics keep evolving -- issue #8 changed the stall
accounting, M2.5 added the AIMD pacer, issue #14 is about to add a wait
policy -- and every round had to be applied three times.

SourceAdmission now owns picking a runnable source and deciding what
happens when none is available. The three loops keep their QuotaGate,
BreakerGate and pacer references because _attempt still needs them for
write-back and pacer.leave(); those instances are shared, not rebuilt
(a second pacer would split the in-flight counter). The cooldown memo
moves in wholesale since only admission consumes it.

Behaviour is unchanged: pick differs from the old chat copy only by the
pacer None-guards, on_no_runnable is verbatim identical, and the suite
reports the same 967 passed / 21 skipped / 32 deselected as before. The
one visible change is the settle-and-release warning text, which had
three variants ("permit", "embedding permit", "OCR permit") and is now
one. Tests importing _demote_call_failures follow it to its new home.
This commit is contained in:
2026-08-19 23:57:01 -04:00
parent 0b3e84b3be
commit 942af99856
5 changed files with 306 additions and 322 deletions
+17 -76
View File
@@ -24,7 +24,6 @@ from loguru import logger
from polygateway.errors import (
AllSourcesExhausted,
CircuitOpenError,
GovernanceBackendError,
PolyGatewayError,
RequestRejectedError,
@@ -33,12 +32,12 @@ from polygateway.errors import (
SourceNotConfiguredError,
TransientError,
)
from polygateway.middleware.admission import SourceAdmission, settle_and_release
from polygateway.middleware.breaker import BreakerGate
from polygateway.middleware.ratelimit import QuotaGate
from polygateway.middleware.retry import StallClock, _failure_reason, backoff_delay
from polygateway.middleware.telemetry import TelemetryEmitter
from polygateway.ports import OutcomeAwareSelector
from polygateway.sources import SourceCooldownMemo
from polygateway.types import (
ChatRequest,
LLMResponse,
@@ -126,14 +125,24 @@ class OcrClient:
self._breaker = BreakerGate(breaker, scope=self._scope)
self._transport = transport
self._retry = retry
self._bp = backpressure
self._quota_full = quota_full
self._emitter = TelemetryEmitter(telemetry, text_cap=text_cap) if telemetry else None
self._telemetry = telemetry
self._memo = SourceCooldownMemo(now=now)
self._now = now
self._sleep = sleep
self._rng = rng
# 准入编排三条循环共用一份(issue #14);冷却备忘由它独占
self._admission = SourceAdmission(
scope=self._scope,
sources=self._sources,
selector=selector,
quota=self._quota,
breaker=self._breaker,
backpressure=backpressure,
quota_full=quota_full,
now=now,
sleep=sleep,
rng=rng,
)
self._closed = False
# —— 公共端口(OcrTextPort / OcrLayoutPort)——
@@ -234,9 +243,9 @@ class OcrClient:
# 只计非生产性等待(issue #8): 真实尝试由重试预算治理,不重复烧 stall 预算
clock = StallClock(self._now)
while True:
picked, gate_rejections = await self._pick_runnable(reasons)
picked, gate_rejections = await self._admission.pick(reasons, {})
if picked is None:
await self._on_no_runnable(gate_rejections, reasons, clock)
await self._admission.on_no_runnable(gate_rejections, reasons, clock)
continue
async with clock.attempting():
outcome = await self._attempt(
@@ -255,62 +264,6 @@ class OcrClient:
if not outcome.immediate:
await self._sleep(backoff_delay(self._retry, fails, outcome.exc, self._rng))
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):
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
entry = None
try:
entry = await self._breaker.try_enter(cand, uuid.uuid4().hex)
finally:
if entry is None:
await self._settle_and_release(permit)
if entry.allowed:
return (cand, permit, entry), gate_rejections
gate_rejections += 1
reasons[cand.name] = "circuit_open"
self._memo.set_until(cand.name, self._now() + entry.retry_after_s)
await self._settle_and_release(permit)
return None, gate_rejections
async def _on_no_runnable(
self, gate_rejections: int, reasons: dict[str, str], clock: StallClock
) -> 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,
)
stall = self._bp.stall_window_s
if clock.stalled_s() > stall and await self._quota.progress_age_s() > stall:
names = tuple(s.name for s in self._sources)
raise AllSourcesExhausted(
scope=self._scope,
reason="stalled",
retry_after_s=await self._breaker.retry_after_s(names),
per_source_reasons=reasons,
)
await self._sleep(self._bp.poll_interval_s * (0.5 + 0.5 * self._rng()))
async def _attempt(
self,
kind: _OcrKind,
@@ -398,7 +351,7 @@ class OcrClient:
)
return _FailedAttempt(exc, immediate=dead)
finally:
await self._settle_and_release(permit)
await settle_and_release(permit, 0)
async def _invoke(
self, kind: _OcrKind, image: bytes, source: SourceConfig, call_id: str
@@ -435,18 +388,6 @@ class OcrClient:
except (GovernanceBackendError, SourceNotConfiguredError) as exc:
logger.warning("OCR 治理记账写回降级(不冒泡): {}", exc)
async def _settle_and_release(self, permit: Permit) -> None:
"""settle 恒 0: OCR 无 token 计费(设计 §5 差异①)。"""
try:
try:
await permit.settle(0)
finally:
await permit.release()
except asyncio.CancelledError:
raise
except Exception as exc:
logger.warning("OCR permit 结算/释放失败(不掩盖主异常): {}", exc)
async def _emit(
self,
kind: _OcrKind,