fix: verify capacity from dual paginated active-task scans
This commit is contained in:
CoolBoy
@@ -325,47 +325,63 @@ class ModelHubClient:
|
||||
return active_count
|
||||
|
||||
def count_active_tasks_exact(self) -> int:
|
||||
"""Conservative active occupancy from filtered, validated pages.
|
||||
"""Confirm active IDs with two independently paginated filtered reads.
|
||||
|
||||
Totals can disagree with records. Count unique IDs, but retain the
|
||||
larger reported total rather than inventing extra free slots.
|
||||
The platform's total/pages fields can describe a different population.
|
||||
Only an explicit empty page terminates a scan, never those fields or a
|
||||
short page. Two page sizes must agree before advertising any free slots.
|
||||
This is not an atomic snapshot: concurrent queue changes fail closed.
|
||||
"""
|
||||
total = 0
|
||||
seen: set[str] = set()
|
||||
details = []
|
||||
for status in ("waiting", "running"):
|
||||
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 invalid(reason: str) -> ModelHubAPIError:
|
||||
error = ModelHubAPIError(reason)
|
||||
error.capacity_reason = reason
|
||||
return error
|
||||
|
||||
def scan(page_size: int) -> tuple[set[str], list[dict[str, Any]]]:
|
||||
active: set[str] = set()
|
||||
details = []
|
||||
for status in ("waiting", "running"):
|
||||
ids: set[str] = set()
|
||||
reported = None
|
||||
for current in range(1, 101):
|
||||
response = self.list_tasks_page(
|
||||
current=current, page_size=page_size, only_mine=True, status=status,
|
||||
)
|
||||
data = response.get("data") if isinstance(response, dict) else None
|
||||
if not isinstance(data, dict) or not isinstance(data.get("records"), list):
|
||||
raise invalid("invalid_active_response")
|
||||
try:
|
||||
value = int(data["total"])
|
||||
if value >= 0:
|
||||
reported = max(reported or 0, value)
|
||||
except (KeyError, ValueError, TypeError):
|
||||
pass # Metadata is diagnostic, not an occupancy source.
|
||||
rows = data["records"]
|
||||
if not rows:
|
||||
break
|
||||
if len(rows) > page_size:
|
||||
raise invalid("page_size_not_respected")
|
||||
previous = len(ids)
|
||||
for row in rows:
|
||||
if not isinstance(row, dict) or not row.get("taskId") or row.get("status") != status:
|
||||
raise invalid("invalid_active_record")
|
||||
ids.add(str(row["taskId"]))
|
||||
if len(ids) == previous:
|
||||
raise invalid("active_pagination_no_progress")
|
||||
else:
|
||||
raise invalid("active_pagination_budget_exhausted")
|
||||
active.update(ids)
|
||||
details.append({"status": status, "reported": reported, "unique": len(ids),
|
||||
"mismatch": reported != len(ids), "pageSize": page_size})
|
||||
return active, details
|
||||
|
||||
self.active_count_details = []
|
||||
first, details = scan(100)
|
||||
confirmed, confirmation_details = scan(50)
|
||||
if first != confirmed:
|
||||
raise invalid("active_snapshot_changed")
|
||||
self.active_count_details = details + confirmation_details
|
||||
return len(confirmed)
|
||||
|
||||
def _request(
|
||||
self,
|
||||
@@ -742,7 +758,8 @@ class ModelHubClientPool:
|
||||
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)
|
||||
reason = getattr(exc, "capacity_reason", "request_failed")
|
||||
print(f"[capacity] refresh_failed error={type(exc).__name__} code={code} reason={reason} count=unknown", flush=True)
|
||||
return None
|
||||
|
||||
def _safe_list_tasks(self, client: ModelHubClient, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
@@ -819,7 +836,7 @@ class ModelHubClientPool:
|
||||
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)
|
||||
f"confirmed_active={new_remote_count} method=dual_page_scan details={details}", flush=True)
|
||||
if new_remote_count > self._account_caps[index]:
|
||||
self._account_caps[index] = new_remote_count
|
||||
caps_changed = True
|
||||
|
||||
Reference in New Issue
Block a user