feat: reserve queue capacity for recent models

modelhub_submmit_api/modelhub_client.py

@@ -9,7 +9,7 @@ from datetime import datetime, timedelta
 from pathlib import Path
 from typing import Any
 
-from common import format_modelhub_datetime, parse_datetime, read_json, runtime_instance_id, write_json
+from common import format_modelhub_datetime, parse_datetime, read_json, runtime_instance_id, utc_now, write_json
 from defaults import EMBEDDED_MODELHUB_XC_TOKEN
 from http_json import HttpJsonError, JsonHttpClient
 
@@ -21,6 +21,10 @@ class ModelHubAPIError(RuntimeError):
         self.payload = payload
 
 
+class OldModelQueuePolicyError(ModelHubAPIError):
+    """No account may accept an older model under the per-account queue policy."""
+
+
 def parse_model_submission_precheck(payload: Any) -> dict[str, Any]:
     """Validate the community lookup response instead of failing open."""
     if not isinstance(payload, dict):
@@ -385,12 +389,26 @@ class ModelHubClientPool:
         capacity_probe_interval_cycles: int = 3,
         capacity_probe_cooldown_cycles: int = 3,
         capacity_state_path: Path | str | None = None,
+        old_model_queue_threshold: int | None = None,
+        recent_model_days: int | None = None,
     ) -> None:
         if not clients:
             raise ValueError("At least one ModelHub client is required")
         self.clients = clients
         configured_cap = active_task_cap if active_task_cap is not None else os.getenv("MODELHUB_AGENT_ACTIVE_TASK_CAP", "100")
         self.active_task_cap = max(1, int(configured_cap))
+        configured_old_threshold = (
+            old_model_queue_threshold
+            if old_model_queue_threshold is not None
+            else os.getenv("MODELHUB_OLD_MODEL_QUEUE_THRESHOLD", "80")
+        )
+        configured_recent_days = (
+            recent_model_days
+            if recent_model_days is not None
+            else os.getenv("MODELHUB_RECENT_MODEL_DAYS", "7")
+        )
+        self.old_model_queue_threshold = max(1, int(configured_old_threshold))
+        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._account_caps = self._load_account_caps(self.active_task_cap)
@@ -671,12 +689,21 @@ class ModelHubClientPool:
     def processed_gpus_for_model(self, model_id: str) -> set[str]:
         return set(self.model_submission_precheck(model_id)["processedGpus"])
 
-    def _reserve_account(self, excluded: set[int]) -> tuple[int, int] | None:
+    def _reserve_account(
+        self,
+        excluded: set[int],
+        *,
+        effective_count_below: int | None = None,
+    ) -> tuple[int, int] | None:
         with self._state_lock:
             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 (
+                    effective_count_below is None
+                    or self._effective_count_locked(index) < effective_count_below
+                )
             }
             usable = {index: remaining for index, remaining in remaining_by_index.items() if remaining > 0}
             if not usable:
@@ -696,7 +723,12 @@ class ModelHubClientPool:
             }
             return selected_index, reservation_id
 
-    def _reserve_probe_account(self, excluded: set[int]) -> tuple[int, int] | None:
+    def _reserve_probe_account(
+        self,
+        excluded: set[int],
+        *,
+        effective_count_below: int | None = None,
+    ) -> tuple[int, int] | None:
         with self._state_lock:
             if not self._capacity_probe_enabled:
                 return None
@@ -707,6 +739,10 @@ class ModelHubClientPool:
                 and index not in self._capacity_probe_attempted
                 and self._capacity_probe_cycle >= self._capacity_probe_cooldown_until[index]
                 and self._effective_count_locked(index) >= self._account_caps[index]
+                and (
+                    effective_count_below is None
+                    or self._effective_count_locked(index) < effective_count_below
+                )
             ]
             if not eligible:
                 return None
@@ -763,17 +799,58 @@ class ModelHubClientPool:
         self._persist_account_caps()
         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(
+                max(
+                    0,
+                    min(self._account_caps[index], self.old_model_queue_threshold)
+                    - self._effective_count_locked(index),
+                )
+                for index in range(len(self.clients))
+            )
+
+    def can_submit_old_models(self) -> bool:
+        return self.old_model_submit_slots() > 0
+
+    def add_task_for_model(
+        self,
+        payload: dict[str, Any],
+        *,
+        model_last_modified: datetime | str | None,
+        submitted_at: datetime | str | None = None,
+    ) -> dict[str, Any]:
+        reference_time = parse_datetime(submitted_at) or utc_now()
+        last_modified = parse_datetime(model_last_modified)
+        is_old_model = (
+            last_modified is None
+            or last_modified < reference_time - timedelta(days=self.recent_model_days)
+        )
+        return self._add_task(payload, old_model_only=is_old_model)
+
     def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
+        return self._add_task(payload, old_model_only=False)
+
+    def _add_task(self, payload: dict[str, Any], *, old_model_only: bool) -> dict[str, Any]:
         self._refresh_active_counts()
         attempted_accounts: set[int] = set()
         forced_refresh_done = False
         last_capacity_error: ModelHubAPIError | None = None
+        account_count_limit = self.old_model_queue_threshold if old_model_only else None
 
         while True:
-            reservation = self._reserve_account(attempted_accounts)
+            reservation = self._reserve_account(
+                attempted_accounts,
+                effective_count_below=account_count_limit,
+            )
             capacity_probe = False
             if reservation is None:
-                reservation = self._reserve_probe_account(attempted_accounts)
+                reservation = self._reserve_probe_account(
+                    attempted_accounts,
+                    effective_count_below=account_count_limit,
+                )
                 capacity_probe = reservation is not None
             if reservation is None:
                 if not forced_refresh_done:
@@ -782,6 +859,11 @@ class ModelHubClientPool:
                     continue
                 if last_capacity_error is not None:
                     raise last_capacity_error
+                if old_model_only:
+                    raise OldModelQueuePolicyError(
+                        f"所有账号活动队列均已达到 {self.old_model_queue_threshold}，"
+                        f"仅允许提交最近 {self.recent_model_days} 天内更新的模型"
+                    )
                 raise ModelHubAPIError(
                     "当前等待中或运行中的异步模型验证任务数量已达已知上限"
                 )