feat: add adaptive GPU scheduling

modelhub_submmit_api/modelhub_client.py

@@ -6,9 +6,10 @@ import threading
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime, timedelta
+from pathlib import Path
 from typing import Any
 
-from common import format_modelhub_datetime, parse_datetime, runtime_instance_id
+from common import format_modelhub_datetime, parse_datetime, read_json, runtime_instance_id, write_json
 from defaults import EMBEDDED_MODELHUB_XC_TOKEN
 from http_json import HttpJsonError, JsonHttpClient
 
@@ -255,6 +256,8 @@ DUPLICATE_SUBMISSION_MARKERS = (
     "already in progress",
 )
 
+DEFAULT_CAPACITY_STATE_PATH = Path(".modelhub_state/account_capacity.json")
+
 
 def is_capacity_error(error: ModelHubAPIError) -> bool:
     if error.code in {409, 429}:
@@ -281,12 +284,24 @@ class ModelHubClientPool:
         active_counts_ttl: float | None = None,
         reservation_ttl: float | None = None,
         instance_id: str | None = None,
+        capacity_probe_interval_cycles: int = 3,
+        capacity_probe_cooldown_cycles: int = 3,
+        capacity_state_path: Path | str | 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))
+        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)
+        self._capacity_probe_interval_cycles = max(0, int(capacity_probe_interval_cycles))
+        self._capacity_probe_cooldown_cycles = max(1, int(capacity_probe_cooldown_cycles))
+        self._capacity_probe_cycle = 0
+        self._capacity_probe_enabled = False
+        self._capacity_probe_attempted: set[int] = set()
+        self._capacity_probe_cooldown_until: list[int] = [0 for _ in clients]
         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
@@ -310,11 +325,65 @@ class ModelHubClientPool:
         # Single reader client to avoid fanout on read operations
         self._reader = clients[0]
 
-    def _safe_count_active_tasks(self, client: ModelHubClient) -> int:
+    @staticmethod
+    def _account_key(client: ModelHubClient, index: int) -> str:
+        token = str(getattr(client, "token", "") or "")
+        identity = token if token else f"account-index:{index}"
+        return hashlib.sha256(identity.encode("utf-8")).hexdigest()[:20]
+
+    def _load_account_caps(self, default_cap: int) -> list[int]:
+        stored: dict[str, Any] = {}
+        if self._capacity_state_path is not None:
+            try:
+                payload = read_json(self._capacity_state_path)
+                if isinstance(payload, dict):
+                    stored = payload.get("accounts") or {}
+            except (FileNotFoundError, ValueError):
+                pass
+        return [
+            max(1, int((stored.get(key) or {}).get("knownCap") or default_cap))
+            for key in self._account_keys
+        ]
+
+    def _persist_account_caps(self) -> None:
+        if self._capacity_state_path is None:
+            return
+        with self._state_lock:
+            payload = {
+                "version": 1,
+                "updatedAt": datetime.now().astimezone().isoformat(),
+                "accounts": {
+                    key: {"accountIndex": index + 1, "knownCap": self._account_caps[index]}
+                    for index, key in enumerate(self._account_keys)
+                },
+            }
+        write_json(self._capacity_state_path, payload)
+
+    def configure_capacity_probe(self, cycle_number: int) -> None:
+        cycle_number = max(0, int(cycle_number))
+        with self._state_lock:
+            if cycle_number != self._capacity_probe_cycle:
+                self._capacity_probe_attempted.clear()
+            self._capacity_probe_cycle = cycle_number
+            self._capacity_probe_enabled = (
+                self._capacity_probe_interval_cycles > 0
+                and cycle_number > 0
+                and cycle_number % self._capacity_probe_interval_cycles == 0
+            )
+
+    def account_capacity_limits(self) -> list[int]:
+        with self._state_lock:
+            return list(self._account_caps)
+
+    def capacity_probe_enabled(self) -> bool:
+        with self._state_lock:
+            return self._capacity_probe_enabled
+
+    def _safe_count_active_tasks(self, client: ModelHubClient, max_count: int) -> int:
         try:
-            return client.count_active_tasks(max_count=self.active_task_cap, page_size=200)
+            return client.count_active_tasks(max_count=max_count, page_size=200)
         except Exception:
-            return self.active_task_cap
+            return max(0, max_count - 1)
 
     def _safe_search_by_model_id(self, client: ModelHubClient, model_id: str) -> dict[str, Any]:
         try:
@@ -346,8 +415,11 @@ class ModelHubClientPool:
                 if not force and self._counts_are_fresh_locked(now):
                     return
 
+            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]:
-                return index, self._safe_count_active_tasks(client)
+                return index, self._safe_count_active_tasks(client, count_limits[index])
 
             results: list[tuple[int, int]] = []
             with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
@@ -363,10 +435,14 @@ class ModelHubClientPool:
                         results.append((index, self.active_task_cap))
 
             refreshed_at = time.monotonic()
+            caps_changed = False
             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)))
+                    new_remote_count = max(0, int(count))
+                    if new_remote_count > self._account_caps[index]:
+                        self._account_caps[index] = new_remote_count
+                        caps_changed = True
                     acknowledged = max(0, new_remote_count - old_remote_count)
                     completed_reservations = sorted(
                         (
@@ -390,12 +466,25 @@ class ModelHubClientPool:
                     self._remote_counts[index] = new_remote_count
                 self._counts_initialized = True
                 self._active_refresh_at = refreshed_at
+            if caps_changed:
+                self._persist_account_caps()
 
     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))]
+        return [self._effective_count_locked(index) for index in range(len(self.clients))]
+
+    def _probe_slot_count_locked(self) -> int:
+        if not self._capacity_probe_enabled:
+            return 0
+        return sum(
+            1
+            for index in range(len(self.clients))
+            if 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]
+        )
 
     def active_task_counts(self) -> list[int]:
         self._refresh_active_counts()
@@ -405,10 +494,11 @@ class ModelHubClientPool:
     def available_submit_slots(self) -> int:
         self._refresh_active_counts()
         with self._state_lock:
-            return sum(
-                max(0, self.active_task_cap - self._effective_count_locked(index))
+            normal_slots = sum(
+                max(0, self._account_caps[index] - self._effective_count_locked(index))
                 for index in range(len(self.clients))
             )
+            return normal_slots if normal_slots > 0 else self._probe_slot_count_locked()
 
     def list_tasks(self, **kwargs):  # noqa: ANN003, ANN001
         # Default: use only the reader client to avoid fanout amplification.
@@ -480,7 +570,7 @@ class ModelHubClientPool:
     def _reserve_account(self, excluded: set[int]) -> tuple[int, int] | None:
         with self._state_lock:
             remaining_by_index = {
-                index: self.active_task_cap - self._effective_count_locked(index)
+                index: self._account_caps[index] - self._effective_count_locked(index)
                 for index in range(len(self.clients))
                 if index not in excluded
             }
@@ -502,6 +592,35 @@ class ModelHubClientPool:
             }
             return selected_index, reservation_id
 
+    def _reserve_probe_account(self, excluded: set[int]) -> tuple[int, int] | None:
+        with self._state_lock:
+            if not self._capacity_probe_enabled:
+                return None
+            eligible = [
+                index
+                for index in range(len(self.clients))
+                if index not in excluded
+                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]
+            ]
+            if not eligible:
+                return None
+            selected_index = min(
+                eligible,
+                key=lambda index: (index - self._selection_cursor) % len(self.clients),
+            )
+            self._selection_cursor = (selected_index + 1) % len(self.clients)
+            self._capacity_probe_attempted.add(selected_index)
+            self._reservation_sequence += 1
+            reservation_id = self._reservation_sequence
+            self._reservations[selected_index][reservation_id] = {
+                "inflight": True,
+                "updated_at": time.monotonic(),
+                "capacity_probe": True,
+            }
+            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)
@@ -513,11 +632,32 @@ class ModelHubClientPool:
             reservation["inflight"] = False
             reservation["updated_at"] = time.monotonic()
 
-    def _mark_account_saturated(self, index: int) -> None:
+    def _mark_account_saturated(self, index: int, *, capacity_probe: bool) -> None:
+        cap_changed = False
         with self._state_lock:
-            self._remote_counts[index] = self.active_task_cap
+            if capacity_probe:
+                self._capacity_probe_cooldown_until[index] = (
+                    self._capacity_probe_cycle + self._capacity_probe_cooldown_cycles
+                )
+            else:
+                observed_capacity = max(1, self._effective_count_locked(index))
+                if observed_capacity < self._account_caps[index]:
+                    self._account_caps[index] = observed_capacity
+                    cap_changed = True
+            self._remote_counts[index] = self._account_caps[index]
             self._counts_initialized = True
             self._active_refresh_at = time.monotonic()
+        if cap_changed:
+            self._persist_account_caps()
+
+    def _promote_account_capacity(self, index: int) -> None:
+        with self._state_lock:
+            discovered_cap = max(self._account_caps[index] + 1, self._effective_count_locked(index))
+            if discovered_cap <= self._account_caps[index]:
+                return
+            self._account_caps[index] = discovered_cap
+        self._persist_account_caps()
+        print(f"[capacity] account={index + 1:02d} discovered_cap={discovered_cap}", flush=True)
 
     def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
         self._refresh_active_counts()
@@ -527,6 +667,10 @@ class ModelHubClientPool:
 
         while True:
             reservation = self._reserve_account(attempted_accounts)
+            capacity_probe = False
+            if reservation is None:
+                reservation = self._reserve_probe_account(attempted_accounts)
+                capacity_probe = reservation is not None
             if reservation is None:
                 if not forced_refresh_done:
                     self._refresh_active_counts(force=True)
@@ -535,7 +679,7 @@ class ModelHubClientPool:
                 if last_capacity_error is not None:
                     raise last_capacity_error
                 raise ModelHubAPIError(
-                    f"当前等待中或运行中的异步模型验证任务数量已达上限（{self.active_task_cap}）"
+                    "当前等待中或运行中的异步模型验证任务数量已达已知上限"
                 )
 
             selected_index, reservation_id = reservation
@@ -548,7 +692,7 @@ class ModelHubClientPool:
                     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)
+                self._mark_account_saturated(selected_index, capacity_probe=capacity_probe)
                 attempted_accounts.add(selected_index)
                 last_capacity_error = exc
                 continue
@@ -557,6 +701,8 @@ class ModelHubClientPool:
                 raise
 
             self._finish_reservation(selected_index, reservation_id, succeeded=True)
+            if capacity_probe:
+                self._promote_account_capacity(selected_index)
             return response
 
     def list_tasks_page(self, **kwargs):  # noqa: ANN003, ANN001