|
|
|
|
@@ -151,7 +151,7 @@ class ModelHubClient:
|
|
|
|
|
"/api/adapt/task/page",
|
|
|
|
|
query={
|
|
|
|
|
"current": current,
|
|
|
|
|
"pageSize": page_size,
|
|
|
|
|
"pageSize": max(1, min(100, int(page_size))),
|
|
|
|
|
"onlyMine": str(only_mine).lower(),
|
|
|
|
|
"beginTime": query_time(begin_time),
|
|
|
|
|
"endTime": query_time(end_time),
|
|
|
|
|
@@ -325,14 +325,46 @@ class ModelHubClient:
|
|
|
|
|
return active_count
|
|
|
|
|
|
|
|
|
|
def count_active_tasks_exact(self) -> int:
|
|
|
|
|
"""Use the platform's filtered totals instead of paging all history."""
|
|
|
|
|
"""Conservative active occupancy from filtered, validated pages.
|
|
|
|
|
|
|
|
|
|
Totals can disagree with records. Count unique IDs, but retain the
|
|
|
|
|
larger reported total rather than inventing extra free slots.
|
|
|
|
|
"""
|
|
|
|
|
total = 0
|
|
|
|
|
seen: set[str] = set()
|
|
|
|
|
details = []
|
|
|
|
|
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"]))
|
|
|
|
|
current, reported = 1, 0
|
|
|
|
|
status_ids: set[str] = set()
|
|
|
|
|
while True:
|
|
|
|
|
data = self.list_tasks_page(current=current, page_size=100, only_mine=True, status=status).get("data")
|
|
|
|
|
if not isinstance(data, dict) or "total" not in data or not isinstance(data.get("records"), list):
|
|
|
|
|
raise ModelHubAPIError(f"Incomplete active-task response for status={status}")
|
|
|
|
|
count = int(data["total"])
|
|
|
|
|
if count < 0:
|
|
|
|
|
raise ModelHubAPIError("Negative active-task total")
|
|
|
|
|
reported = max(reported, count)
|
|
|
|
|
pages = max(1, int(data.get("pages") or 0), (count + 99) // 100)
|
|
|
|
|
rows = data["records"]
|
|
|
|
|
prior = len(status_ids)
|
|
|
|
|
for row in rows:
|
|
|
|
|
if not isinstance(row, dict) or not row.get("taskId") or row.get("status") != status:
|
|
|
|
|
raise ModelHubAPIError(f"Invalid filtered active-task record for status={status}")
|
|
|
|
|
status_ids.add(str(row["taskId"]))
|
|
|
|
|
if current >= pages:
|
|
|
|
|
if current > 1 and len(status_ids) == prior:
|
|
|
|
|
raise ModelHubAPIError("Active-task pagination made no progress")
|
|
|
|
|
break
|
|
|
|
|
if not rows or len(status_ids) == prior or current >= 1000:
|
|
|
|
|
raise ModelHubAPIError("Incomplete active-task pagination")
|
|
|
|
|
current += 1
|
|
|
|
|
# An ID can move from waiting to running between requests.
|
|
|
|
|
overlap = len(status_ids & seen)
|
|
|
|
|
total += max(reported, len(status_ids)) - overlap
|
|
|
|
|
seen.update(status_ids)
|
|
|
|
|
details.append({"status": status, "reported": reported, "unique": len(status_ids),
|
|
|
|
|
"mismatch": reported != len(status_ids)})
|
|
|
|
|
self.active_count_details = details
|
|
|
|
|
return total
|
|
|
|
|
|
|
|
|
|
def _request(
|
|
|
|
|
@@ -521,6 +553,7 @@ class ModelHubClientPool:
|
|
|
|
|
# persisted capacity is not enough without a successful active-count
|
|
|
|
|
# read for the current scheduling window.
|
|
|
|
|
self._count_known: list[bool] = [False for _ in clients]
|
|
|
|
|
self._count_success_at = [0.0 for _ in clients]
|
|
|
|
|
self._active_refresh_at: float = 0.0
|
|
|
|
|
self._counts_initialized = False
|
|
|
|
|
self._state_lock = threading.Lock()
|
|
|
|
|
@@ -703,10 +736,13 @@ class ModelHubClientPool:
|
|
|
|
|
|
|
|
|
|
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"):
|
|
|
|
|
if 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 client.count_active_tasks(max_count=max_count, page_size=100)
|
|
|
|
|
except Exception as exc:
|
|
|
|
|
# Never log request URLs/headers or credentials.
|
|
|
|
|
code = getattr(exc, "code", None)
|
|
|
|
|
print(f"[capacity] refresh_failed error={type(exc).__name__} code={code} count=unknown", flush=True)
|
|
|
|
|
return None
|
|
|
|
|
|
|
|
|
|
def _safe_list_tasks(self, client: ModelHubClient, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
|
|
|
|
|
@@ -779,6 +815,11 @@ class ModelHubClientPool:
|
|
|
|
|
old_remote_count = self._remote_counts[index]
|
|
|
|
|
new_remote_count = max(0, int(count))
|
|
|
|
|
self._count_known[index] = True
|
|
|
|
|
self._count_success_at[index] = refreshed_at
|
|
|
|
|
details = getattr(self.clients[index], "active_count_details", [])
|
|
|
|
|
if any(item.get("mismatch") for item in details):
|
|
|
|
|
print(f"[capacity] account={index + 1:02d} total_records_mismatch=true "
|
|
|
|
|
f"conservative_active={new_remote_count} details={details}", flush=True)
|
|
|
|
|
if new_remote_count > self._account_caps[index]:
|
|
|
|
|
self._account_caps[index] = new_remote_count
|
|
|
|
|
caps_changed = True
|
|
|
|
|
@@ -821,6 +862,7 @@ class ModelHubClientPool:
|
|
|
|
|
1
|
|
|
|
|
for index in range(len(self.clients))
|
|
|
|
|
if index not in self._capacity_probe_attempted
|
|
|
|
|
and self._count_known[index]
|
|
|
|
|
and self._capacity_probe_cycle >= self._capacity_probe_cooldown_until[index]
|
|
|
|
|
and self._effective_count_locked(index) >= self._account_caps[index]
|
|
|
|
|
)
|
|
|
|
|
@@ -846,9 +888,20 @@ class ModelHubClientPool:
|
|
|
|
|
normal_slots = sum(
|
|
|
|
|
max(0, self._account_caps[index] - self._effective_count_locked(index))
|
|
|
|
|
for index in range(len(self.clients))
|
|
|
|
|
if self._count_known[index]
|
|
|
|
|
)
|
|
|
|
|
return normal_slots if normal_slots > 0 else self._probe_slot_count_locked()
|
|
|
|
|
|
|
|
|
|
def active_count_status(self) -> list[dict[str, Any]]:
|
|
|
|
|
self._refresh_active_counts()
|
|
|
|
|
with self._state_lock:
|
|
|
|
|
now = time.monotonic()
|
|
|
|
|
return [{"account": i + 1, "known": self._count_known[i],
|
|
|
|
|
"active": self._effective_count_locked(i) if self._count_known[i] else None,
|
|
|
|
|
"lastSuccessAgeSeconds": round(now - self._count_success_at[i], 1)
|
|
|
|
|
if self._count_success_at[i] else None}
|
|
|
|
|
for i in range(len(self.clients))]
|
|
|
|
|
|
|
|
|
|
def list_tasks(self, **kwargs): # noqa: ANN003, ANN001
|
|
|
|
|
# Default: use only the reader client to avoid fanout amplification.
|
|
|
|
|
# The fanout_all parameter allows explicit cross-account merging when needed.
|
|
|
|
|
@@ -935,6 +988,7 @@ class ModelHubClientPool:
|
|
|
|
|
index: self._account_caps[index] - self._effective_count_locked(index)
|
|
|
|
|
for index in range(len(self.clients))
|
|
|
|
|
if index not in excluded
|
|
|
|
|
and self._count_known[index]
|
|
|
|
|
and (weekly_target is None or index == weekly_target)
|
|
|
|
|
and (
|
|
|
|
|
reserve_capacity_slots is None
|
|
|
|
|
@@ -997,6 +1051,7 @@ class ModelHubClientPool:
|
|
|
|
|
index
|
|
|
|
|
for index in range(len(self.clients))
|
|
|
|
|
if index not in excluded
|
|
|
|
|
and self._count_known[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 (
|
|
|
|
|
|