fix: coordinate concurrent account capacity filling

modelhub_submmit_api/modelhub_client.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
+import hashlib
 import os
-import time
 import threading
+import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime, timedelta
 from typing import Any
 
-from common import format_modelhub_datetime, parse_datetime
+from common import format_modelhub_datetime, parse_datetime, runtime_instance_id
 from defaults import EMBEDDED_MODELHUB_XC_TOKEN
 from http_json import HttpJsonError, JsonHttpClient
 
@@ -233,18 +234,63 @@ def is_active_task(task: dict[str, Any]) -> bool:
     return True
 
 
+CAPACITY_ERROR_MARKERS = (
+    "达到上限",
+    "达上限",
+    "任务数量已达",
+    "队列已满",
+    "queue is full",
+    "queue full",
+    "capacity",
+    "too many active",
+    "active task limit",
+)
+
+
+def is_capacity_error(error: ModelHubAPIError) -> bool:
+    if error.code in {409, 429}:
+        return True
+    message = str(error).strip().lower()
+    if isinstance(error.payload, dict):
+        message = f"{message} {error.payload.get('message') or ''}".lower()
+    return any(marker in message for marker in CAPACITY_ERROR_MARKERS)
+
+
 class ModelHubClientPool:
-    def __init__(self, clients: list[ModelHubClient], *, active_task_cap: int = 100) -> None:
+    def __init__(
+        self,
+        clients: list[ModelHubClient],
+        *,
+        active_task_cap: int | None = None,
+        active_counts_ttl: float | None = None,
+        reservation_ttl: float | None = None,
+        instance_id: str | None = None,
+    ) -> None:
         if not clients:
             raise ValueError("At least one ModelHub client is required")
         self.clients = clients
-        self.active_task_cap = active_task_cap
-        self._active_counts: list[int] = [0 for _ in 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_ttl = active_counts_ttl if active_counts_ttl is not None else os.getenv("MODELHUB_AGENT_ACTIVE_COUNTS_TTL_SECONDS", "15")
+        configured_reservation_ttl = (
+            reservation_ttl
+            if reservation_ttl is not None
+            else os.getenv("MODELHUB_AGENT_RESERVATION_TTL_SECONDS", "120")
+        )
+        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]
         self._active_refresh_at: float = 0.0
-        self._active_counts_ttl: float = 30.0
+        self._counts_initialized = False
         self._state_lock = threading.Lock()
-        # Per-cycle cache for search_by_model_id results (model_id -> merged verify result map)
-        self._verify_cache: dict[str, dict[str, Any]] = {}
+        self._refresh_lock = threading.Lock()
+        self._reservations: list[dict[int, dict[str, Any]]] = [{} for _ in clients]
+        self._reservation_sequence = 0
+        identity = instance_id or runtime_instance_id()
+        identity_hash = hashlib.sha256(identity.encode("utf-8")).hexdigest()
+        self._selection_cursor = int(identity_hash[:12], 16) % len(clients)
+        self._verify_cache: dict[str, tuple[float, dict[str, Any]]] = {}
+        self._verify_cache_ttl = max(1.0, float(os.getenv("MODELHUB_AGENT_VERIFY_CACHE_TTL_SECONDS", "30")))
         # Single reader client to avoid fanout on read operations
         self._reader = clients[0]
 
@@ -267,48 +313,86 @@ class ModelHubClientPool:
         except Exception:
             return []
 
+    def _counts_are_fresh_locked(self, now: float) -> bool:
+        return self._counts_initialized and (now - self._active_refresh_at) < self._active_counts_ttl
+
     def _refresh_active_counts(self, *, force: bool = False) -> None:
-        now = time.time()
-        if (not force) and self._active_counts and (now - self._active_refresh_at) < self._active_counts_ttl:
-            return
+        now = time.monotonic()
+        with self._state_lock:
+            if not force and self._counts_are_fresh_locked(now):
+                return
 
-        def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int]:
-            return index, self._safe_count_active_tasks(client)
+        # Do not hold the scheduler lock during remote I/O. One refresher is
+        # enough; all submitting threads can continue using their reservations.
+        with self._refresh_lock:
+            now = time.monotonic()
+            with self._state_lock:
+                if not force and self._counts_are_fresh_locked(now):
+                    return
 
-        results: list[tuple[int, int]] = []
-        with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
-            futures = {executor.submit(_to_indexed_result, index, client): index for index, client in enumerate(self.clients)}
-            for future in as_completed(futures):
-                index = futures[future]
-                try:
-                    results.append(future.result())
-                except Exception:
-                    results.append((index, self.active_task_cap))
+            def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int]:
+                return index, self._safe_count_active_tasks(client)
 
-        self._active_counts = [0 for _ in self.clients]
-        for index, count in sorted(results, key=lambda item: item[0]):
-            self._active_counts[index] = count
-        # Use current time after refresh completes, not the stale 'now' from function start
-        self._active_refresh_at = time.time()
+            results: list[tuple[int, int]] = []
+            with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
+                futures = {
+                    executor.submit(_to_indexed_result, index, client): index
+                    for index, client in enumerate(self.clients)
+                }
+                for future in as_completed(futures):
+                    index = futures[future]
+                    try:
+                        results.append(future.result())
+                    except Exception:
+                        results.append((index, self.active_task_cap))
+
+            refreshed_at = time.monotonic()
+            with self._state_lock:
+                for index, count in results:
+                    old_remote_count = self._remote_counts[index]
+                    new_remote_count = min(self.active_task_cap, max(0, int(count)))
+                    acknowledged = max(0, new_remote_count - old_remote_count)
+                    completed_reservations = sorted(
+                        (
+                            (reservation_id, reservation)
+                            for reservation_id, reservation in self._reservations[index].items()
+                            if not reservation["inflight"]
+                        ),
+                        key=lambda item: item[1]["updated_at"],
+                    )
+                    for reservation_id, _reservation in completed_reservations[:acknowledged]:
+                        self._reservations[index].pop(reservation_id, None)
+
+                    # A remote count can stay flat when one old task finishes as
+                    # one new task appears. Expiry prevents that net-zero update
+                    # from reserving a slot forever.
+                    for reservation_id, reservation in list(self._reservations[index].items()):
+                        if reservation["inflight"]:
+                            continue
+                        if refreshed_at - reservation["updated_at"] >= self._reservation_ttl:
+                            self._reservations[index].pop(reservation_id, None)
+                    self._remote_counts[index] = new_remote_count
+                self._counts_initialized = True
+                self._active_refresh_at = refreshed_at
+
+    def _effective_count_locked(self, index: int) -> int:
+        return self._remote_counts[index] + len(self._reservations[index])
+
+    def _counts_snapshot_locked(self) -> list[int]:
+        return [min(self.active_task_cap, self._effective_count_locked(index)) for index in range(len(self.clients))]
 
     def active_task_counts(self) -> list[int]:
+        self._refresh_active_counts()
         with self._state_lock:
-            now = time.time()
-            if self._active_counts and (now - self._active_refresh_at) < self._active_counts_ttl:
-                pass  # use cached counts
-            else:
-                self._refresh_active_counts()
-        return list(self._active_counts)
+            return self._counts_snapshot_locked()
 
     def available_submit_slots(self) -> int:
+        self._refresh_active_counts()
         with self._state_lock:
-            # Only refresh if cache is stale (respects TTL) or counts are empty
-            now = time.time()
-            if self._active_counts and (now - self._active_refresh_at) < self._active_counts_ttl:
-                pass  # use cached counts
-            else:
-                self._refresh_active_counts()
-            return sum(max(0, self.active_task_cap - count) for count in self._active_counts)
+            return sum(
+                max(0, self.active_task_cap - self._effective_count_locked(index))
+                for index in range(len(self.clients))
+            )
 
     def list_tasks(self, **kwargs):  # noqa: ANN003, ANN001
         # Default: use only the reader client to avoid fanout amplification.
@@ -339,10 +423,19 @@ class ModelHubClientPool:
         cached = self._verify_cache.get(model_id)
         if cached is None:
             return None
-        return dict(cached)
+        cached_at, payload = cached
+        if time.monotonic() - cached_at >= self._verify_cache_ttl:
+            self._verify_cache.pop(model_id, None)
+            return None
+        return dict(payload)
 
     def _write_to_cache(self, model_id: str, payload: dict[str, Any]) -> None:
-        self._verify_cache[model_id] = payload
+        self._verify_cache[model_id] = (time.monotonic(), payload)
+
+    def begin_cycle(self) -> None:
+        """Drop model verification cache entries from the previous scan cycle."""
+        with self._state_lock:
+            self._verify_cache.clear()
 
     def search_by_model_id(self, model_id: str) -> dict[str, Any]:
         # Check cache first (per-cycle cache to avoid repeated API calls for the same model)
@@ -368,30 +461,87 @@ class ModelHubClientPool:
     def processed_gpus_for_model(self, model_id: str) -> set[str]:
         return set(self.get_verify_result_map(model_id).keys())
 
-    def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
+    def _reserve_account(self, excluded: set[int]) -> tuple[int, int] | None:
         with self._state_lock:
-            self._refresh_active_counts(force=True)
-            selected_index = None
-            best_remaining = -1
-            for index, active_count in enumerate(self._active_counts):
-                remaining = self.active_task_cap - active_count
-                if remaining > best_remaining:
-                    best_remaining = remaining
-                    selected_index = index
-            if selected_index is None or best_remaining <= 0:
+            remaining_by_index = {
+                index: self.active_task_cap - self._effective_count_locked(index)
+                for index in range(len(self.clients))
+                if index not in excluded
+            }
+            usable = {index: remaining for index, remaining in remaining_by_index.items() if remaining > 0}
+            if not usable:
+                return None
+            best_remaining = max(usable.values())
+            tied = [index for index, remaining in usable.items() if remaining == best_remaining]
+            selected_index = min(
+                tied,
+                key=lambda index: (index - self._selection_cursor) % len(self.clients),
+            )
+            self._selection_cursor = (selected_index + 1) % len(self.clients)
+            self._reservation_sequence += 1
+            reservation_id = self._reservation_sequence
+            self._reservations[selected_index][reservation_id] = {
+                "inflight": True,
+                "updated_at": time.monotonic(),
+            }
+            return selected_index, reservation_id
+
+    def _finish_reservation(self, index: int, reservation_id: int, *, succeeded: bool) -> None:
+        with self._state_lock:
+            reservation = self._reservations[index].get(reservation_id)
+            if reservation is None:
+                return
+            if not succeeded:
+                self._reservations[index].pop(reservation_id, None)
+                return
+            reservation["inflight"] = False
+            reservation["updated_at"] = time.monotonic()
+
+    def _mark_account_saturated(self, index: int) -> None:
+        with self._state_lock:
+            self._remote_counts[index] = self.active_task_cap
+            self._counts_initialized = True
+            self._active_refresh_at = time.monotonic()
+
+    def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
+        self._refresh_active_counts()
+        attempted_accounts: set[int] = set()
+        forced_refresh_done = False
+        last_capacity_error: ModelHubAPIError | None = None
+
+        while True:
+            reservation = self._reserve_account(attempted_accounts)
+            if reservation is None:
+                if not forced_refresh_done:
+                    self._refresh_active_counts(force=True)
+                    forced_refresh_done = True
+                    continue
+                if last_capacity_error is not None:
+                    raise last_capacity_error
                 raise ModelHubAPIError(
                     f"当前等待中或运行中的异步模型验证任务数量已达上限（{self.active_task_cap}）"
                 )
-            self._active_counts[selected_index] += 1
 
-        selected_client = self.clients[selected_index]
-        try:
-            response = selected_client.add_task(payload)
-        except Exception:
-            with self._state_lock:
-                self._active_counts[selected_index] -= 1
-            raise
-        return response
+            selected_index, reservation_id = reservation
+            selected_client = self.clients[selected_index]
+            try:
+                response = selected_client.add_task(payload)
+            except ModelHubAPIError as exc:
+                self._finish_reservation(selected_index, reservation_id, succeeded=False)
+                if not is_capacity_error(exc):
+                    raise
+                # Another process may have filled this account after our count
+                # refresh. Mark it full locally and immediately try another one.
+                self._mark_account_saturated(selected_index)
+                attempted_accounts.add(selected_index)
+                last_capacity_error = exc
+                continue
+            except Exception:
+                self._finish_reservation(selected_index, reservation_id, succeeded=False)
+                raise
+
+            self._finish_reservation(selected_index, reservation_id, succeeded=True)
+            return response
 
     def list_tasks_page(self, **kwargs):  # noqa: ANN003, ANN001
         """Single-page task listing via the reader client (no fanout)."""