feat: make model age admission-only

modelhub_submmit_api/modelhub_client.py

@@ -407,7 +407,7 @@ class ModelHubClientPool:
         configured_recent_reserve = (
             recent_model_reserve_slots
             if recent_model_reserve_slots is not None
-            else os.getenv("MODELHUB_RECENT_MODEL_RESERVE_SLOTS", "10")
+            else os.getenv("MODELHUB_RECENT_MODEL_RESERVE_SLOTS", "5")
         )
         configured_recent_days = (
             recent_model_days
@@ -434,6 +434,10 @@ class ModelHubClientPool:
         self._active_counts_ttl = max(1.0, float(configured_ttl))
         self._reservation_ttl = max(self._active_counts_ttl * 2, float(configured_reservation_ttl))
         self._remote_counts: list[int] = [0 for _ in clients]
+        # Old-model admission is fail-closed per account. A configured or
+        # 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._active_refresh_at: float = 0.0
         self._counts_initialized = False
         self._state_lock = threading.Lock()
@@ -518,11 +522,11 @@ class ModelHubClientPool:
         with self._state_lock:
             return self._capacity_probe_enabled
 
-    def _safe_count_active_tasks(self, client: ModelHubClient, max_count: int) -> int:
+    def _safe_count_active_tasks(self, client: ModelHubClient, max_count: int) -> int | None:
         try:
             return client.count_active_tasks(max_count=max_count, page_size=200)
         except Exception:
-            return max(0, max_count - 1)
+            return None
 
     def _safe_list_tasks(self, client: ModelHubClient, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
         try:
@@ -550,10 +554,10 @@ class ModelHubClientPool:
             with self._state_lock:
                 count_limits = [cap + 1 for cap in self._account_caps]
 
-            def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int]:
+            def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int | None]:
                 return index, self._safe_count_active_tasks(client, count_limits[index])
 
-            results: list[tuple[int, int]] = []
+            results: list[tuple[int, int | None]] = []
             with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
                 futures = {
                     executor.submit(_to_indexed_result, index, client): index
@@ -564,14 +568,18 @@ class ModelHubClientPool:
                     try:
                         results.append(future.result())
                     except Exception:
-                        results.append((index, self.active_task_cap))
+                        results.append((index, None))
 
             refreshed_at = time.monotonic()
             caps_changed = False
             with self._state_lock:
                 for index, count in results:
+                    if count is None:
+                        self._count_known[index] = False
+                        continue
                     old_remote_count = self._remote_counts[index]
                     new_remote_count = max(0, int(count))
+                    self._count_known[index] = True
                     if new_remote_count > self._account_caps[index]:
                         self._account_caps[index] = new_remote_count
                         caps_changed = True
@@ -725,8 +733,11 @@ class ModelHubClientPool:
                 if index not in excluded
                 and (
                     reserve_capacity_slots is None
-                    or self._effective_count_locked(index)
-                    < max(0, self._account_caps[index] - reserve_capacity_slots)
+                    or (
+                        self._count_known[index]
+                        and 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}
@@ -765,8 +776,11 @@ class ModelHubClientPool:
                 and self._effective_count_locked(index) >= self._account_caps[index]
                 and (
                     reserve_capacity_slots is None
-                    or self._effective_count_locked(index)
-                    < max(0, self._account_caps[index] - reserve_capacity_slots)
+                    or (
+                        self._count_known[index]
+                        and self._effective_count_locked(index)
+                        < max(0, self._account_caps[index] - reserve_capacity_slots)
+                    )
                 )
             ]
             if not eligible:
@@ -810,6 +824,7 @@ class ModelHubClientPool:
                     self._account_caps[index] = observed_capacity
                     cap_changed = True
             self._remote_counts[index] = self._account_caps[index]
+            self._count_known[index] = True
             self._counts_initialized = True
             self._active_refresh_at = time.monotonic()
         if cap_changed:
@@ -834,15 +849,19 @@ class ModelHubClientPool:
                     self._old_model_count_limit_locked(index) - self._effective_count_locked(index),
                 )
                 for index in range(len(self.clients))
+                if self._count_known[index]
             )
 
     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."""
+    def old_model_queue_thresholds(self) -> list[int | None]:
+        """Return thresholds, or None when old-model admission must fail closed."""
         with self._state_lock:
-            return [self._old_model_count_limit_locked(index) for index in range(len(self.clients))]
+            return [
+                self._old_model_count_limit_locked(index) if self._count_known[index] else None
+                for index in range(len(self.clients))
+            ]
 
     def can_submit_old_models(self) -> bool:
         return self.old_model_submit_slots() > 0