diff --git a/src/polygateway/embedding.py b/src/polygateway/embedding.py index 55e501a..d7e9b70 100644 --- a/src/polygateway/embedding.py +++ b/src/polygateway/embedding.py @@ -39,7 +39,7 @@ from polygateway.errors import ( ) from polygateway.middleware.breaker import BreakerGate from polygateway.middleware.ratelimit import QuotaGate -from polygateway.middleware.retry import _failure_reason, backoff_delay +from polygateway.middleware.retry import StallClock, _failure_reason, backoff_delay from polygateway.middleware.telemetry import TelemetryEmitter from polygateway.sources import SourceCooldownMemo from polygateway.types import ( @@ -179,13 +179,15 @@ class EmbeddingClient: ) -> _BatchOutcome: fails = 0 reasons: dict[str, str] = {} - entered_at = self._now() + # 只计非生产性等待(issue #8): 真实尝试由重试预算治理,不重复烧 stall 预算 + clock = StallClock(self._now) while True: picked, gate_rejections = await self._pick_runnable(reasons) if picked is None: - await self._on_no_runnable(gate_rejections, reasons, entered_at) + await self._on_no_runnable(gate_rejections, reasons, clock) continue - outcome = await self._attempt(batch, *picked, reasons, session_id, parent_call_id) + async with clock.attempting(): + outcome = await self._attempt(batch, *picked, reasons, session_id, parent_call_id) if isinstance(outcome, _BatchOutcome): return outcome fails += 1 @@ -228,7 +230,7 @@ class EmbeddingClient: return None, gate_rejections async def _on_no_runnable( - self, gate_rejections: int, reasons: dict[str, str], entered_at: float + 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) @@ -245,7 +247,7 @@ class EmbeddingClient: per_source_reasons=reasons, ) stall = self._bp.stall_window_s - if self._now() - entered_at > stall and await self._quota.progress_age_s() > stall: + 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, diff --git a/tests/unit/test_embedding.py b/tests/unit/test_embedding.py index a33f5a8..8519583 100644 --- a/tests/unit/test_embedding.py +++ b/tests/unit/test_embedding.py @@ -173,6 +173,8 @@ from polygateway.types import ( # noqa: E402 GlobalLimits, RetryPolicy, ) +from tests.contracts.conftest import FakeClock # noqa: E402 +from tests.unit.test_backpressure import BoundedSleep # noqa: E402 _BREAKER = BreakerConfig(fail_threshold=3, cooldown_s=60.0, probe_ttl_s=120.0) _NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0) @@ -208,6 +210,31 @@ class ScriptedEmbedTransport: return action +class _ClockAdvancingEmbedTransport: + """按脚本 [(推进秒数, 动作), ...] 执行: 在一次尝试内部推进时钟(issue #8)。 + + 动作语义同 `ScriptedEmbedTransport`。stall 口径要区分"时间花在哪", + 故必须能让时钟只在 transport 内前进。 + """ + + def __init__(self, script, clock): + self.script = list(script) + self.clock = clock + self.calls = [] + + async def embed(self, *, texts, source, call_id): + self.calls.append((source.name, list(texts), call_id)) + advance, action = self.script.pop(0) + self.clock.advance(advance) + if isinstance(action, Exception): + raise action + if action == "hang": + await asyncio.Event().wait() + if action == "ok": + return _vec_for(texts) + return action + + class _MemoryRecorder: def __init__(self): self.rows = [] @@ -338,6 +365,39 @@ class TestEmbedGovernance: await task assert (await limiter.source_stats("e1")).inflight == 0 + async def test_single_timeout_does_not_exhaust_stall_budget(self): + """issue #8: 一次耗满 timeout 的尝试不得吃掉 stall 预算而使重试失效。 + + embedding 只有 `_on_no_runnable` 一处 stall 判定, 故失效链条是 + "先超时一次(墙钟耗尽) → 再遇到无可用源 → 判死"。此处正是这条路径。 + """ + clock = FakeClock() + limiter = held = None # 闭包延迟求值: client 建好后才有 limiter + + async def toggle_permit(n): + """首次退避占满 permit, 迫使下一轮走 _on_no_runnable; 之后放行。""" + nonlocal held + if n == 1: + held = await limiter.try_acquire("e1", 0) + else: + await held.release() + + # 第一次尝试耗满 300s 超时失败, 随后被迫走一轮 _on_no_runnable—— + # stall 判定就在那里, 检验它有没有把这 300s 生产性时间算进 stall 账 + transport = _ClockAdvancingEmbedTransport( + [(300.1, TransientError("timeout", status_code=504)), (0.0, "ok")], clock + ) + client, limiter = _embed_client( + [_src(max_concurrency=1)], + [], + now=clock, + transport=transport, + sleep=BoundedSleep(toggle_permit), + ) + resp = await client.embed(["a"]) + assert resp.vectors == [[1.0]] + assert len(transport.calls) == 2 # 第二次尝试确实发出了 + class TestEmbedTelemetry: async def test_per_batch_rows_with_digest(self):