fix: respect task page limits and recover stale capacity counts

modelhub_submmit_api/modelhub_client.py

@@ -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 (