Token load balancer with load-aware selection
| 15 | |
| 16 | |
| 17 | class LoadBalancer: |
| 18 | """Token load balancer with load-aware selection""" |
| 19 | |
| 20 | def __init__(self, token_manager, concurrency_manager: Optional[ConcurrencyManager] = None): |
| 21 | self.token_manager = token_manager |
| 22 | self.concurrency_manager = concurrency_manager |
| 23 | self._image_pending: Dict[int, int] = {} |
| 24 | self._video_pending: Dict[int, int] = {} |
| 25 | self._pending_lock = asyncio.Lock() |
| 26 | self._round_robin_state: Dict[str, Optional[int]] = {"image": None, "video": None, "default": None} |
| 27 | self._rr_lock = asyncio.Lock() |
| 28 | |
| 29 | async def _get_pending_count(self, token_id: int, for_image_generation: bool, for_video_generation: bool) -> int: |
| 30 | async with self._pending_lock: |
| 31 | if for_image_generation: |
| 32 | return max(0, int(self._image_pending.get(token_id, 0))) |
| 33 | if for_video_generation: |
| 34 | return max(0, int(self._video_pending.get(token_id, 0))) |
| 35 | return 0 |
| 36 | |
| 37 | async def _add_pending(self, token_id: int, for_image_generation: bool, for_video_generation: bool): |
| 38 | async with self._pending_lock: |
| 39 | if for_image_generation: |
| 40 | self._image_pending[token_id] = max(0, int(self._image_pending.get(token_id, 0))) + 1 |
| 41 | elif for_video_generation: |
| 42 | self._video_pending[token_id] = max(0, int(self._video_pending.get(token_id, 0))) + 1 |
| 43 | |
| 44 | async def release_pending(self, token_id: int, for_image_generation: bool = False, for_video_generation: bool = False): |
| 45 | async with self._pending_lock: |
| 46 | if for_image_generation: |
| 47 | current = max(0, int(self._image_pending.get(token_id, 0))) |
| 48 | if current <= 1: |
| 49 | self._image_pending.pop(token_id, None) |
| 50 | else: |
| 51 | self._image_pending[token_id] = current - 1 |
| 52 | elif for_video_generation: |
| 53 | current = max(0, int(self._video_pending.get(token_id, 0))) |
| 54 | if current <= 1: |
| 55 | self._video_pending.pop(token_id, None) |
| 56 | else: |
| 57 | self._video_pending[token_id] = current - 1 |
| 58 | |
| 59 | async def _get_token_load(self, token_id: int, for_image_generation: bool, for_video_generation: bool) -> tuple[int, Optional[int]]: |
| 60 | """获取 token 当前负载。 |
| 61 | |
| 62 | Returns: |
| 63 | (inflight, remaining) |
| 64 | remaining 为 None 表示无限制 |
| 65 | """ |
| 66 | if not self.concurrency_manager: |
| 67 | return 0, None |
| 68 | |
| 69 | if for_image_generation: |
| 70 | inflight = await self.concurrency_manager.get_image_inflight(token_id) |
| 71 | remaining = await self.concurrency_manager.get_image_remaining(token_id) |
| 72 | pending = await self._get_pending_count(token_id, True, False) |
| 73 | effective_inflight = inflight + pending |
| 74 | if remaining is not None: |