feat: derive recent queue reserve from account capacity

modelhub_submmit_api/queue_cleanup.py

@@ -242,12 +242,11 @@ def find_old_overflow_tasks(
     tasks: list[OwnedTask],
     *,
     model_last_modified: dict[str, datetime],
-    queue_threshold: int = 80,
+    queue_threshold: int | dict[int, int] = 80,
     recent_model_days: int = 7,
     reference_time: datetime | None = None,
     incomplete_accounts: set[int] | None = None,
 ) -> tuple[list[dict[str, Any]], dict[str, int]]:
-    threshold = max(1, int(queue_threshold))
     recent_days = max(1, int(recent_model_days))
     cutoff = (reference_time or utc_now()) - timedelta(days=recent_days)
     incomplete_accounts = incomplete_accounts or set()
@@ -261,10 +260,19 @@ def find_old_overflow_tasks(
         "withinFirstQueuePositions": 0,
         "recentOverflowTasks": 0,
         "modelAgeUnknown": 0,
+        "accountThresholdUnknown": 0,
     }
     for account_index, account_tasks in sorted(grouped.items()):
         if account_index in incomplete_accounts:
             continue
+        if isinstance(queue_threshold, dict):
+            configured_threshold = queue_threshold.get(account_index)
+            if configured_threshold is None:
+                skipped["accountThresholdUnknown"] += 1
+                continue
+        else:
+            configured_threshold = queue_threshold
+        threshold = max(0, int(configured_threshold))
         ordered = sorted(account_tasks, key=lambda item: item.task_id)
         skipped["withinFirstQueuePositions"] += min(threshold, len(ordered))
         for position, task in enumerate(ordered, start=1):
@@ -307,19 +315,19 @@ def cleanup_certain_oom_tasks(
     read_concurrency: int = 6,
     stop_batch_size: int = DEFAULT_STOP_BATCH_SIZE,
     gpu_memory_gib: dict[str, float] | None = None,
-    age_queue_threshold: int | None = None,
+    age_reserved_slots: int | None = None,
     reference_time: datetime | None = None,
     log: Callable[[str], None] = print,
 ) -> dict[str, Any]:
     """Stop deterministic OOM tasks and old tasks beyond each account's protected prefix."""
     clients = list(modelhub.clients)
     reference_time = reference_time or utc_now()
-    configured_queue_threshold = (
-        age_queue_threshold
-        if age_queue_threshold is not None
-        else getattr(modelhub, "old_model_queue_threshold", 80)
+    configured_reserved_slots = (
+        age_reserved_slots
+        if age_reserved_slots is not None
+        else getattr(modelhub, "recent_model_reserve_slots", 10)
     )
-    queue_threshold = max(1, int(configured_queue_threshold or 80))
+    reserved_slots = max(0, int(configured_reserved_slots or 0))
     recent_model_days = max(1, int(getattr(modelhub, "recent_model_days", 7) or 7))
     tasks, listing_errors = collect_active_tasks(clients, read_concurrency=read_concurrency)
     log(
@@ -327,6 +335,25 @@ def cleanup_certain_oom_tasks(
         f"listing_errors={sum(len(items) for items in listing_errors.values())}"
     )
 
+    observed_active_counts: list[int | None] = [0 for _ in clients]
+    for task in tasks:
+        current_count = observed_active_counts[task.account_index]
+        if current_count is not None:
+            observed_active_counts[task.account_index] = current_count + 1
+    for account_index in listing_errors:
+        observed_active_counts[account_index] = None
+    if hasattr(modelhub, "observe_capacity_lower_bounds"):
+        account_caps = list(modelhub.observe_capacity_lower_bounds(observed_active_counts))
+    elif hasattr(modelhub, "account_capacity_limits"):
+        account_caps = list(modelhub.account_capacity_limits())
+    else:
+        default_cap = max(1, int(getattr(modelhub, "active_task_cap", 100) or 100))
+        account_caps = [default_cap for _ in clients]
+    queue_thresholds = {
+        index: max(0, int(account_caps[index]) - reserved_slots)
+        for index in range(min(len(clients), len(account_caps)))
+    }
+
     model_ids = {task.model_id for task in tasks}
     repository_sizes, size_errors = _load_repository_sizes(
         model_ids,
@@ -362,7 +389,7 @@ def cleanup_certain_oom_tasks(
         for task in sorted(
             (item for item in age_rank_tasks if item.account_index == account_index),
             key=lambda item: item.task_id,
-        )[queue_threshold:]
+        )[queue_thresholds.get(account_index, len(age_rank_tasks)):]
     }
     model_last_modified, age_errors = _load_model_last_modified(
         overflow_model_ids,
@@ -373,14 +400,15 @@ def cleanup_certain_oom_tasks(
     old_overflow_decisions, age_skipped = find_old_overflow_tasks(
         age_rank_tasks,
         model_last_modified=model_last_modified,
-        queue_threshold=queue_threshold,
+        queue_threshold=queue_thresholds,
         recent_model_days=recent_model_days,
         reference_time=reference_time,
         incomplete_accounts=set(listing_errors),
     )
     log(
         f"[queue-cleanup] old_overflow={len(old_overflow_decisions)} "
-        f"threshold={queue_threshold} recent_days={recent_model_days} "
+        f"thresholds={','.join(str(queue_thresholds[index]) for index in sorted(queue_thresholds))} "
+        f"reserve_recent_slots={reserved_slots} recent_days={recent_model_days} "
         f"recent_overflow={age_skipped['recentOverflowTasks']} "
         f"age_unknown={age_skipped['modelAgeUnknown']}"
     )
@@ -453,7 +481,12 @@ def cleanup_certain_oom_tasks(
             cleanup_reasons = set(decision.get("cleanupReasons") or [decision.get("reason")])
             age_only = cleanup_reasons == {"old_model_beyond_account_queue_threshold"}
             current_position = active_positions_by_account.get(account_index, {}).get(int(decision["taskId"]))
-            if age_only and (current_position is None or current_position <= queue_threshold):
+            account_queue_threshold = queue_thresholds.get(account_index)
+            if age_only and (
+                current_position is None
+                or account_queue_threshold is None
+                or current_position <= account_queue_threshold
+            ):
                 policy_no_longer_applies.append(
                     {**decision, "recheckedQueuePosition": current_position}
                 )
@@ -516,7 +549,14 @@ def cleanup_certain_oom_tasks(
         "certainOomTasks": oom_decisions,
         "oldOverflowCount": len(old_overflow_decisions),
         "oldOverflowTasks": old_overflow_decisions,
-        "oldModelQueueThreshold": queue_threshold,
+        "oldModelQueueThresholds": [
+            queue_thresholds.get(index) for index in range(len(clients))
+        ],
+        "accountCapacityLimits": [
+            account_caps[index] if index < len(account_caps) else None
+            for index in range(len(clients))
+        ],
+        "recentModelReserveSlots": reserved_slots,
         "recentModelDays": recent_model_days,
         "agePolicySkipped": age_skipped,
         "cleanupCandidateCount": len(decisions),