feat: calibrate account capacity weekly

This commit is contained in:
CoolBoy
2026-09-20 13:03:35 +08:00
parent 75c034d588
commit 4afce879dc
7 changed files with 466 additions and 38 deletions

View File

@@ -8,6 +8,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo
from common import format_modelhub_datetime, parse_datetime, read_json, runtime_instance_id, utc_now, write_json
from defaults import EMBEDDED_MODELHUB_XC_TOKEN
@@ -310,6 +311,17 @@ class ModelHubClient:
return active_count
def count_active_tasks_exact(self) -> int:
"""Use the platform's filtered totals instead of paging all history."""
total = 0
for status in ("waiting", "running"):
page = self.list_tasks_page(current=1, page_size=1, only_mine=True, status=status)
data = page.get("data")
if not isinstance(data, dict) or "total" not in data:
raise ModelHubAPIError(f"Active-task total is unavailable for status={status}")
total += max(0, int(data["total"]))
return total
def _request(
self,
method: str,
@@ -409,8 +421,6 @@ DEFAULT_CAPACITY_STATE_PATH = Path(".modelhub_state/account_capacity.json")
def is_capacity_error(error: ModelHubAPIError) -> bool:
if error.code in {409, 429}:
return True
message = str(error).strip().lower()
if isinstance(error.payload, dict):
message = f"{message} {error.payload.get('message') or ''}".lower()
@@ -447,7 +457,7 @@ class ModelHubClientPool:
active_counts_ttl: float | None = None,
reservation_ttl: float | None = None,
instance_id: str | None = None,
capacity_probe_interval_cycles: int = 3,
capacity_probe_interval_cycles: int = 0,
capacity_probe_cooldown_cycles: int = 3,
capacity_state_path: Path | str | None = None,
recent_model_reserve_slots: int | None = None,
@@ -472,7 +482,13 @@ class ModelHubClientPool:
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)]
self._weekly_calibration_weeks: list[str | None] = [None for _ in clients]
self._weekly_calibration_results: list[str | None] = [None for _ in clients]
self._account_caps = self._load_account_caps(self.active_task_cap)
self._weekly_calibration_week: str | None = None
self._weekly_calibration_window = False
self._weekly_calibration_deadline = 0.0
self._weekly_no_candidate_retry_at = 0.0
self._capacity_probe_interval_cycles = max(0, int(capacity_probe_interval_cycles))
self._capacity_probe_cooldown_cycles = max(1, int(capacity_probe_cooldown_cycles))
self._capacity_probe_cycle = 0
@@ -521,10 +537,15 @@ class ModelHubClientPool:
stored = payload.get("accounts") or {}
except (FileNotFoundError, ValueError):
pass
return [
max(1, int((stored.get(key) or {}).get("knownCap") or default_cap))
for key in self._account_keys
]
caps: list[int] = []
for index, key in enumerate(self._account_keys):
account = stored.get(key) or {}
if not isinstance(account, dict):
account = {}
caps.append(max(1, int(account.get("knownCap") or default_cap)))
self._weekly_calibration_weeks[index] = account.get("calibrationWeek") or None
self._weekly_calibration_results[index] = account.get("calibrationResult") or None
return caps
def _persist_account_caps(self) -> None:
if self._capacity_state_path is None:
@@ -534,7 +555,12 @@ class ModelHubClientPool:
"version": 1,
"updatedAt": datetime.now().astimezone().isoformat(),
"accounts": {
key: {"accountIndex": index + 1, "knownCap": self._account_caps[index]}
key: {
"accountIndex": index + 1,
"knownCap": self._account_caps[index],
"calibrationWeek": self._weekly_calibration_weeks[index],
"calibrationResult": self._weekly_calibration_results[index],
}
for index, key in enumerate(self._account_keys)
},
}
@@ -552,6 +578,92 @@ class ModelHubClientPool:
and cycle_number % self._capacity_probe_interval_cycles == 0
)
def configure_weekly_calibration(self, now: datetime) -> dict[str, Any]:
"""Only Monday 10:00-10:59 in Beijing may borrow recent-model slots."""
local_now = now.astimezone(ZoneInfo("Asia/Shanghai"))
week = f"{local_now.isocalendar().year}-W{local_now.isocalendar().week:02d}"
window = local_now.weekday() == 0 and local_now.hour == 10
with self._state_lock:
if week != self._weekly_calibration_week:
self._weekly_no_candidate_retry_at = 0.0
self._weekly_calibration_week = week if window else None
self._weekly_calibration_window = window
self._weekly_calibration_deadline = (
time.monotonic()
+ (local_now.replace(hour=11, minute=0, second=0, microsecond=0) - local_now).total_seconds()
if window else 0.0
)
target = self._weekly_target_locked()
return {
"window": window,
"week": week,
"targetAccount": target + 1 if target is not None else None,
"completedAccounts": sum(value == week for value in self._weekly_calibration_weeks),
}
def _weekly_target_locked(self) -> int | None:
if (
not self._weekly_calibration_window
or not self._weekly_calibration_week
or time.monotonic() >= self._weekly_calibration_deadline
or time.monotonic() < self._weekly_no_candidate_retry_at
):
return None
for index in range(len(self.clients)):
if (
self._weekly_calibration_weeks[index] != self._weekly_calibration_week
and self._count_known[index]
):
return index
return None
def weekly_calibration_active(self) -> bool:
self._refresh_active_counts()
with self._state_lock:
return self._weekly_target_locked() is not None
def weekly_calibration_status(self) -> dict[str, Any]:
self._refresh_active_counts()
with self._state_lock:
target = self._weekly_target_locked()
return {
"window": self._weekly_calibration_window and time.monotonic() < self._weekly_calibration_deadline,
"week": self._weekly_calibration_week,
"targetAccount": target + 1 if target is not None else None,
"completedAccounts": sum(
value == self._weekly_calibration_week
for value in self._weekly_calibration_weeks
) if self._weekly_calibration_week else 0,
}
def finish_weekly_calibration(self, index: int, *, result: str) -> None:
with self._state_lock:
if not self._weekly_calibration_window or not self._weekly_calibration_week:
return
self._weekly_calibration_weeks[index] = self._weekly_calibration_week
self._weekly_calibration_results[index] = result
known_cap = self._account_caps[index]
active = self._effective_count_locked(index)
threshold = self._old_model_count_limit_locked(index)
self._persist_account_caps()
print(
f"[capacity-calibration] account={index + 1:02d} result={result} "
f"known_cap={known_cap} active={active} old_threshold={threshold}",
flush=True,
)
def pause_weekly_calibration_for_no_candidates(self) -> None:
with self._state_lock:
target = self._weekly_target_locked()
if target is None:
return
self._weekly_no_candidate_retry_at = time.monotonic() + 600
print(
f"[capacity-calibration] account={target + 1:02d} "
"status=paused reason=no_candidates retry_seconds=600",
flush=True,
)
def account_capacity_limits(self) -> list[int]:
with self._state_lock:
return list(self._account_caps)
@@ -576,8 +688,10 @@ class ModelHubClientPool:
with self._state_lock:
return self._capacity_probe_enabled
def _safe_count_active_tasks(self, client: ModelHubClient, max_count: int) -> int | None:
def _safe_count_active_tasks(self, client: ModelHubClient, max_count: int | None) -> int | None:
try:
if max_count is None and hasattr(client, "count_active_tasks_exact"):
return client.count_active_tasks_exact()
return client.count_active_tasks(max_count=max_count, page_size=200)
except Exception:
return None
@@ -606,7 +720,25 @@ class ModelHubClientPool:
return
with self._state_lock:
count_limits = [cap + 1 for cap in self._account_caps]
count_limits: list[int | None] = [cap + 1 for cap in self._account_caps]
weekly_target = self._weekly_target_locked()
if (
weekly_target is None
and self._weekly_calibration_window
and time.monotonic() < self._weekly_calibration_deadline
and time.monotonic() >= self._weekly_no_candidate_retry_at
):
weekly_target = next(
(
index for index, week in enumerate(self._weekly_calibration_weeks)
if week != self._weekly_calibration_week
),
None,
)
if weekly_target is not None:
# A cap+1 early stop only proves a lower bound. Weekly
# calibration needs an exact active count for its target.
count_limits[weekly_target] = None
def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int | None]:
return index, self._safe_count_active_tasks(client, count_limits[index])
@@ -694,6 +826,10 @@ class ModelHubClientPool:
def available_submit_slots(self) -> int:
self._refresh_active_counts()
with self._state_lock:
if self._weekly_target_locked() is not None:
# The weekly run may need to cross the previously known cap.
# Bound each scan; it can resume in the next cycle.
return 100
normal_slots = sum(
max(0, self._account_caps[index] - self._effective_count_locked(index))
for index in range(len(self.clients))
@@ -781,10 +917,12 @@ class ModelHubClientPool:
reserve_capacity_slots: int | None = None,
) -> tuple[int, int] | None:
with self._state_lock:
weekly_target = self._weekly_target_locked()
remaining_by_index = {
index: self._account_caps[index] - self._effective_count_locked(index)
for index in range(len(self.clients))
if index not in excluded
and (weekly_target is None or index == weekly_target)
and (
reserve_capacity_slots is None
or (
@@ -819,14 +957,19 @@ class ModelHubClientPool:
reserve_capacity_slots: int | None = None,
) -> tuple[int, int] | None:
with self._state_lock:
if not self._capacity_probe_enabled:
weekly_target = self._weekly_target_locked()
if weekly_target is None and not self._capacity_probe_enabled:
return None
eligible = [
index
for index in range(len(self.clients))
if index not in excluded
and index not in self._capacity_probe_attempted
and self._capacity_probe_cycle >= self._capacity_probe_cooldown_until[index]
and (weekly_target is None or index == weekly_target)
and (weekly_target is not None or index not in self._capacity_probe_attempted)
and (
weekly_target is not None
or self._capacity_probe_cycle >= self._capacity_probe_cooldown_until[index]
)
and self._effective_count_locked(index) >= self._account_caps[index]
and (
reserve_capacity_slots is None
@@ -844,7 +987,8 @@ class ModelHubClientPool:
key=lambda index: (index - self._selection_cursor) % len(self.clients),
)
self._selection_cursor = (selected_index + 1) % len(self.clients)
self._capacity_probe_attempted.add(selected_index)
if weekly_target is None:
self._capacity_probe_attempted.add(selected_index)
self._reservation_sequence += 1
reservation_id = self._reservation_sequence
self._reservations[selected_index][reservation_id] = {
@@ -877,7 +1021,12 @@ class ModelHubClientPool:
if observed_capacity < self._account_caps[index]:
self._account_caps[index] = observed_capacity
cap_changed = True
self._remote_counts[index] = self._account_caps[index]
# Accepted local reservations may not have appeared in the last
# remote count yet. Do not count them twice when marking full.
self._remote_counts[index] = max(
0,
self._account_caps[index] - len(self._reservations[index]),
)
self._count_known[index] = True
self._counts_initialized = True
self._active_refresh_at = time.monotonic()
@@ -890,14 +1039,16 @@ class ModelHubClientPool:
if discovered_cap <= self._account_caps[index]:
return
self._account_caps[index] = discovered_cap
weekly = self._weekly_target_locked() == index
self._persist_account_caps()
print(f"[capacity] account={index + 1:02d} discovered_cap={discovered_cap}", flush=True)
if not weekly or discovered_cap % 25 == 0:
print(f"[capacity] account={index + 1:02d} discovered_cap={discovered_cap}", flush=True)
def old_model_submit_slots(self) -> int:
"""Return queue positions where models older than the recent window are allowed."""
self._refresh_active_counts()
with self._state_lock:
return sum(
normal_slots = sum(
max(
0,
self._old_model_count_limit_locked(index) - self._effective_count_locked(index),
@@ -905,6 +1056,7 @@ class ModelHubClientPool:
for index in range(len(self.clients))
if self._count_known[index]
)
return max(1, normal_slots) if self._weekly_target_locked() is not None else normal_slots
def _old_model_count_limit_locked(self, index: int) -> int:
return max(0, self._account_caps[index] - self.recent_model_reserve_slots)
@@ -943,9 +1095,15 @@ class ModelHubClientPool:
attempted_accounts: set[int] = set()
forced_refresh_done = False
last_capacity_error: ModelHubAPIError | None = None
reserved_capacity_slots = self.recent_model_reserve_slots if old_model_only else None
while True:
with self._state_lock:
weekly_target = self._weekly_target_locked()
reserved_capacity_slots = (
self.recent_model_reserve_slots
if old_model_only and weekly_target is None
else None
)
reservation = self._reserve_account(
attempted_accounts,
reserve_capacity_slots=reserved_capacity_slots,
@@ -977,6 +1135,9 @@ class ModelHubClientPool:
selected_index, reservation_id = reservation
selected_client = self.clients[selected_index]
if weekly_target == selected_index and time.monotonic() >= self._weekly_calibration_deadline:
self._finish_reservation(selected_index, reservation_id, succeeded=False)
continue
try:
response = selected_client.add_task(payload)
except ModelHubAPIError as exc:
@@ -986,6 +1147,18 @@ class ModelHubClientPool:
# Another process may have filled this account after our count
# refresh. Mark it full locally and immediately try another one.
self._mark_account_saturated(selected_index, capacity_probe=capacity_probe)
if weekly_target == selected_index:
self._refresh_active_counts(force=True)
with self._state_lock:
confirmed = (
self._count_known[selected_index]
and self._remote_counts[selected_index] >= self._account_caps[selected_index]
)
self.finish_weekly_calibration(
selected_index,
result="saturated" if confirmed else "inconclusive_capacity_response",
)
self._refresh_active_counts(force=True)
attempted_accounts.add(selected_index)
last_capacity_error = exc
continue