feat: add source selectors and cooldown memo
This commit is contained in:
@@ -0,0 +1,61 @@
|
|||||||
|
"""选源策略与源冷却备忘(M1 设计 §2.3;蓝本 CHS app/providers/selector.py 与 governance.py:107)。
|
||||||
|
|
||||||
|
选源是端口(SourceSelector),两个首发实现逐字移植 CHS;冷却备忘是
|
||||||
|
RetryMW 的进程本地状态——熔断开路的源在本地记冷却截止,选源时跳过,
|
||||||
|
避免每轮白烧 RPM 去探测已知开路的源。
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import time
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Callable
|
||||||
|
|
||||||
|
from polygateway.types import SourceConfig, SourceStats
|
||||||
|
|
||||||
|
|
||||||
|
class RoundRobinSelector:
|
||||||
|
"""轮转起点后移(CHS selector.py:20 同款);单 client 内游标推进。"""
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self._n = 0
|
||||||
|
|
||||||
|
def order(
|
||||||
|
self, sources: list[SourceConfig], stats: dict[str, SourceStats]
|
||||||
|
) -> list[SourceConfig]:
|
||||||
|
if not sources:
|
||||||
|
return []
|
||||||
|
k = self._n % len(sources)
|
||||||
|
self._n += 1
|
||||||
|
return sources[k:] + sources[:k]
|
||||||
|
|
||||||
|
|
||||||
|
class LeastInflightSelector:
|
||||||
|
"""最少在途优先(CHS selector.py:36 同款);排序稳定,平局保持配置序。"""
|
||||||
|
|
||||||
|
def order(
|
||||||
|
self, sources: list[SourceConfig], stats: dict[str, SourceStats]
|
||||||
|
) -> list[SourceConfig]:
|
||||||
|
return sorted(sources, key=lambda s: stats[s.name].inflight if s.name in stats else 0)
|
||||||
|
|
||||||
|
|
||||||
|
class SourceCooldownMemo:
|
||||||
|
"""进程本地的源冷却备忘(CHS governance.py:107 同款)。
|
||||||
|
|
||||||
|
只记"冷却截止时刻";set_until 取更晚者,防止较早的提示回退已有备忘。
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, now: Callable[[], float] = time.monotonic) -> None:
|
||||||
|
self._now = now
|
||||||
|
self._until: dict[str, float] = {}
|
||||||
|
|
||||||
|
def set_until(self, source_name: str, until: float) -> None:
|
||||||
|
self._until[source_name] = max(self._until.get(source_name, 0.0), until)
|
||||||
|
|
||||||
|
def active(self, source_name: str) -> bool:
|
||||||
|
return self._until.get(source_name, 0.0) > self._now()
|
||||||
|
|
||||||
|
def remaining(self, source_name: str) -> float:
|
||||||
|
return max(0.0, self._until.get(source_name, 0.0) - self._now())
|
||||||
@@ -0,0 +1,58 @@
|
|||||||
|
"""sources.py 选源策略与冷却备忘测试(蓝本 CHS app/providers/selector.py:20,36)。"""
|
||||||
|
|
||||||
|
from polygateway.sources import LeastInflightSelector, RoundRobinSelector, SourceCooldownMemo
|
||||||
|
from polygateway.types import SourceConfig, SourceStats
|
||||||
|
|
||||||
|
|
||||||
|
def _src(name):
|
||||||
|
return SourceConfig(
|
||||||
|
name=name, provider="openai", base_url="https://gw.example/v1",
|
||||||
|
api_key="sk", model="m", timeout_s=10.0,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestRoundRobin:
|
||||||
|
def test_rotation_advances_each_call(self):
|
||||||
|
sources = [_src("a"), _src("b"), _src("c")]
|
||||||
|
sel = RoundRobinSelector()
|
||||||
|
orders = [tuple(s.name for s in sel.order(sources, {})) for _ in range(4)]
|
||||||
|
assert orders == [("a", "b", "c"), ("b", "c", "a"), ("c", "a", "b"), ("a", "b", "c")]
|
||||||
|
|
||||||
|
def test_single_source_stable(self):
|
||||||
|
sources = [_src("only")]
|
||||||
|
sel = RoundRobinSelector()
|
||||||
|
assert [s.name for s in sel.order(sources, {})] == ["only"]
|
||||||
|
assert [s.name for s in sel.order(sources, {})] == ["only"]
|
||||||
|
|
||||||
|
|
||||||
|
class TestLeastInflight:
|
||||||
|
def test_orders_by_inflight_ascending(self):
|
||||||
|
sources = [_src("a"), _src("b"), _src("c")]
|
||||||
|
stats = {"a": SourceStats(5, 0, 0), "b": SourceStats(1, 0, 0), "c": SourceStats(3, 0, 0)}
|
||||||
|
sel = LeastInflightSelector()
|
||||||
|
assert [s.name for s in sel.order(sources, stats)] == ["b", "c", "a"]
|
||||||
|
|
||||||
|
def test_stable_on_ties(self):
|
||||||
|
sources = [_src("a"), _src("b")]
|
||||||
|
stats = {"a": SourceStats(2, 0, 0), "b": SourceStats(2, 0, 0)}
|
||||||
|
assert [s.name for s in LeastInflightSelector().order(sources, stats)] == ["a", "b"]
|
||||||
|
|
||||||
|
|
||||||
|
class TestCooldownMemo:
|
||||||
|
def test_active_until_expiry(self):
|
||||||
|
t = {"now": 100.0}
|
||||||
|
memo = SourceCooldownMemo(now=lambda: t["now"])
|
||||||
|
assert not memo.active("a")
|
||||||
|
memo.set_until("a", 130.0)
|
||||||
|
assert memo.active("a")
|
||||||
|
assert memo.remaining("a") == 30.0
|
||||||
|
t["now"] = 131.0
|
||||||
|
assert not memo.active("a")
|
||||||
|
assert memo.remaining("a") == 0.0
|
||||||
|
|
||||||
|
def test_extend_after_wait_hint(self):
|
||||||
|
t = {"now": 0.0}
|
||||||
|
memo = SourceCooldownMemo(now=lambda: t["now"])
|
||||||
|
memo.set_until("a", 10.0)
|
||||||
|
memo.set_until("a", 5.0) # 更早的截止不回退已有备忘
|
||||||
|
assert memo.remaining("a") == 10.0
|
||||||
Reference in New Issue
Block a user