165 lines
5.1 KiB
Python
165 lines
5.1 KiB
Python
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)
|