初始化提交

2026-08-21 16:00:13 +08:00
commit 887c1f449e
@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import random
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+
+@dataclass
+class PoolContext:
+    request_id: str = ""
+    excluded_endpoints: set[str] = field(default_factory=set)
+
+
+class SchedulingStrategy(Protocol):
+    def select(self, excluded: set[str]) -> dict[str, Any] | None: ...
+
+
+class WeightedRoundRobinStrategy:
+    """策略 1：在所有活跃端点间加权轮询。"""
+
+    def __init__(self, endpoints: list[dict[str, Any]]) -> None:
+        self._endpoints = endpoints
+        self._index = 0
+
+    def select(self, excluded: set[str]) -> dict[str, Any] | None:
+        available = [
+            e
+            for e in self._endpoints
+            if e.get("name") not in excluded and e.get("status", "active") == "active"
+        ]
+        if not available:
+            return None
+        weights = [e.get("weight", 100) for e in available]
+        self._index = (self._index + 1) % len(available)
+        return random.choices(available, weights=weights, k=1)[0]
+
+
+class PriorityFailoverStrategy:
+    """策略 2：选择最高优先级；当排除时顺延。"""
+
+    def __init__(self, endpoints: list[dict[str, Any]]) -> None:
+        self._endpoints = sorted(endpoints, key=lambda e: e.get("priority", 100))
+
+    def select(self, excluded: set[str]) -> dict[str, Any] | None:
+        for ep in self._endpoints:
+            if ep.get("name") not in excluded and ep.get("status", "active") == "active":
+                return ep
+        return None
+
+
+class LowestLatencyStrategy:
+    """策略 3：选择最近延迟最低的端点。"""
+
+    def __init__(self, endpoints: list[dict[str, Any]]) -> None:
+        self._endpoints = endpoints
+
+    def select(self, excluded: set[str]) -> dict[str, Any] | None:
+        available = [
+            e
+            for e in self._endpoints
+            if e.get("name") not in excluded and e.get("status", "active") == "active"
+        ]
+        if not available:
+            return None
+        return min(available, key=lambda e: e.get("latency_ms", 1000))
+
+
+class LeastConnectionsStrategy:
+    """策略 4：选择活跃连接数最少的端点。"""
+
+    def __init__(self, endpoints: list[dict[str, Any]]) -> None:
+        self._endpoints = endpoints
+
+    def select(self, excluded: set[str]) -> dict[str, Any] | None:
+        available = [
+            e
+            for e in self._endpoints
+            if e.get("name") not in excluded and e.get("status", "active") == "active"
+        ]
+        if not available:
+            return None
+        return min(available, key=lambda e: e.get("connections", 0))
+
+
+class CostAwareStrategy:
+    """策略 5：按每 token 成本选择最便宜的端点。"""
+
+    def __init__(
+        self, endpoints: list[dict[str, Any]], get_cost: Callable[[str], float] | None = None
+    ) -> None:
+        self._endpoints = endpoints
+        self._get_cost = get_cost or (lambda _: 0.0)
+
+    def select(self, excluded: set[str]) -> dict[str, Any] | None:
+        available = [
+            e
+            for e in self._endpoints
+            if e.get("name") not in excluded and e.get("status", "active") == "active"
+        ]
+        if not available:
+            return None
+        return min(available, key=lambda e: self._get_cost(e.get("name", "")))
+
+
+_STRATEGIES = {
+    "weighted_round_robin": WeightedRoundRobinStrategy,
+    "priority_failover": PriorityFailoverStrategy,
+    "lowest_latency": LowestLatencyStrategy,
+    "least_connections": LeastConnectionsStrategy,
+    "cost_aware": CostAwareStrategy,
+}
+
+
+class EndpointPool:
+    """模型池，带可插拔的调度策略。"""
+
+    def __init__(
+        self,
+        logical_model_name: str,
+        endpoints_configs: list[dict[str, Any]],
+        strategy: str = "weighted_round_robin",
+    ) -> None:
+        self.logical_model_name = logical_model_name
+        self._endpoints = endpoints_configs
+        strategy_cls = _STRATEGIES.get(strategy, WeightedRoundRobinStrategy)
+        self.strategy = (
+            strategy_cls(endpoints_configs)
+            if strategy == "priority_failover"
+            else strategy_cls(endpoints_configs)
+        )
+
+    def select(self, excluded: set[str]) -> dict[str, Any] | None:
+        """使用配置的策略选择一个端点。"""
+        return self.strategy.select(excluded)
+
+    def list_endpoints(self) -> list[dict[str, Any]]:
+        return list(self._endpoints)
+
+
+class PoolFactory:
+    """工厂：创建端点池。"""
+
+    @staticmethod
+    def create_pool(
+        logical_model_name: str,
+        endpoints: list[dict[str, Any]],
+        strategy_name: str = "weighted_round_robin",
+    ) -> EndpointPool:
+        return EndpointPool(
+            logical_model_name=logical_model_name,
+            endpoints_configs=endpoints,
+            strategy=strategy_name,
+        )
+
+    @staticmethod
+    def create_strategy(
+        name: str,
+        endpoints: list[dict[str, Any]],
+    ) -> SchedulingStrategy:
+        strategy_cls = _STRATEGIES.get(name)
+        if strategy_cls is None:
+            raise ValueError(f"Unknown scheduling strategy: {name}")
+        return strategy_cls(endpoints)