feat: derive recent queue reserve from account capacity

This commit is contained in:
CoolBoy
2026-08-11 01:41:26 +08:00
parent f12d96b138
commit 2065ad6abc
11 changed files with 266 additions and 132 deletions

View File

@@ -389,7 +389,7 @@ class ModelHubClientPool:
capacity_probe_interval_cycles: int = 3,
capacity_probe_cooldown_cycles: int = 3,
capacity_state_path: Path | str | None = None,
old_model_queue_threshold: int | None = None,
recent_model_reserve_slots: int | None = None,
recent_model_days: int | None = None,
) -> None:
if not clients:
@@ -397,17 +397,17 @@ class ModelHubClientPool:
self.clients = clients
configured_cap = active_task_cap if active_task_cap is not None else os.getenv("MODELHUB_AGENT_ACTIVE_TASK_CAP", "100")
self.active_task_cap = max(1, int(configured_cap))
configured_old_threshold = (
old_model_queue_threshold
if old_model_queue_threshold is not None
else os.getenv("MODELHUB_OLD_MODEL_QUEUE_THRESHOLD", "80")
configured_recent_reserve = (
recent_model_reserve_slots
if recent_model_reserve_slots is not None
else os.getenv("MODELHUB_RECENT_MODEL_RESERVE_SLOTS", "10")
)
configured_recent_days = (
recent_model_days
if recent_model_days is not None
else os.getenv("MODELHUB_RECENT_MODEL_DAYS", "7")
)
self.old_model_queue_threshold = max(1, int(configured_old_threshold))
self.recent_model_reserve_slots = max(0, int(configured_recent_reserve))
self.recent_model_days = max(1, int(configured_recent_days))
self._capacity_state_path = Path(capacity_state_path) if capacity_state_path else None
self._account_keys = [self._account_key(client, index) for index, client in enumerate(clients)]
@@ -491,6 +491,22 @@ class ModelHubClientPool:
with self._state_lock:
return list(self._account_caps)
def observe_capacity_lower_bounds(self, active_counts: list[int | None]) -> list[int]:
"""Promote known caps from complete account listings without guessing an upper bound."""
caps_changed = False
with self._state_lock:
for index, value in enumerate(active_counts[: len(self._account_caps)]):
if value is None:
continue
observed_count = max(0, int(value))
if observed_count > self._account_caps[index]:
self._account_caps[index] = observed_count
caps_changed = True
result = list(self._account_caps)
if caps_changed:
self._persist_account_caps()
return result
def capacity_probe_enabled(self) -> bool:
with self._state_lock:
return self._capacity_probe_enabled
@@ -693,7 +709,7 @@ class ModelHubClientPool:
self,
excluded: set[int],
*,
effective_count_below: int | None = None,
reserve_capacity_slots: int | None = None,
) -> tuple[int, int] | None:
with self._state_lock:
remaining_by_index = {
@@ -701,8 +717,9 @@ class ModelHubClientPool:
for index in range(len(self.clients))
if index not in excluded
and (
effective_count_below is None
or self._effective_count_locked(index) < effective_count_below
reserve_capacity_slots is None
or self._effective_count_locked(index)
< max(0, self._account_caps[index] - reserve_capacity_slots)
)
}
usable = {index: remaining for index, remaining in remaining_by_index.items() if remaining > 0}
@@ -727,7 +744,7 @@ class ModelHubClientPool:
self,
excluded: set[int],
*,
effective_count_below: int | None = None,
reserve_capacity_slots: int | None = None,
) -> tuple[int, int] | None:
with self._state_lock:
if not self._capacity_probe_enabled:
@@ -740,8 +757,9 @@ class ModelHubClientPool:
and self._capacity_probe_cycle >= self._capacity_probe_cooldown_until[index]
and self._effective_count_locked(index) >= self._account_caps[index]
and (
effective_count_below is None
or self._effective_count_locked(index) < effective_count_below
reserve_capacity_slots is None
or self._effective_count_locked(index)
< max(0, self._account_caps[index] - reserve_capacity_slots)
)
]
if not eligible:
@@ -806,12 +824,19 @@ class ModelHubClientPool:
return sum(
max(
0,
min(self._account_caps[index], self.old_model_queue_threshold)
- self._effective_count_locked(index),
self._old_model_count_limit_locked(index) - self._effective_count_locked(index),
)
for index in range(len(self.clients))
)
def _old_model_count_limit_locked(self, index: int) -> int:
return max(0, self._account_caps[index] - self.recent_model_reserve_slots)
def old_model_queue_thresholds(self) -> list[int]:
"""Return each account's current cap minus its recent-model reserve."""
with self._state_lock:
return [self._old_model_count_limit_locked(index) for index in range(len(self.clients))]
def can_submit_old_models(self) -> bool:
return self.old_model_submit_slots() > 0
@@ -838,18 +863,18 @@ class ModelHubClientPool:
attempted_accounts: set[int] = set()
forced_refresh_done = False
last_capacity_error: ModelHubAPIError | None = None
account_count_limit = self.old_model_queue_threshold if old_model_only else None
reserved_capacity_slots = self.recent_model_reserve_slots if old_model_only else None
while True:
reservation = self._reserve_account(
attempted_accounts,
effective_count_below=account_count_limit,
reserve_capacity_slots=reserved_capacity_slots,
)
capacity_probe = False
if reservation is None:
reservation = self._reserve_probe_account(
attempted_accounts,
effective_count_below=account_count_limit,
reserve_capacity_slots=reserved_capacity_slots,
)
capacity_probe = reservation is not None
if reservation is None:
@@ -860,8 +885,10 @@ class ModelHubClientPool:
if last_capacity_error is not None:
raise last_capacity_error
if old_model_only:
thresholds = self.old_model_queue_thresholds()
raise OldModelQueuePolicyError(
f"所有账号活动队列均已达到 {self.old_model_queue_threshold},"
f"所有账号均已达到动态旧模型阈值 {thresholds}(账号上限减 "
f"{self.recent_model_reserve_slots}),"
f"仅允许提交最近 {self.recent_model_days} 天内更新的模型"
)
raise ModelHubAPIError(