Switch agent discovery to ModelScope

modelhub_submmit_api/modelhub_client.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import os
 from datetime import datetime, timedelta
+from itertools import cycle
+from threading import Lock
 from typing import Any
 
 from common import format_modelhub_datetime, parse_datetime
@@ -205,6 +207,73 @@ class ModelHubClient:
         return payload
 
 
+class MultiModelHubClient:
+    def __init__(
+        self,
+        *,
+        tokens: list[str],
+        base_url: str = "https://modelhub.org.cn",
+        timeout: int = 30,
+        retries: int = 2,
+    ) -> None:
+        deduped: list[str] = []
+        for token in tokens:
+            clean = token.strip()
+            if clean and clean not in deduped:
+                deduped.append(clean)
+        if not deduped:
+            raise ValueError("At least one ModelHub token is required")
+
+        self.clients = [
+            ModelHubClient(token=token, base_url=base_url, timeout=timeout, retries=retries)
+            for token in deduped
+        ]
+        self._cycle = cycle(self.clients)
+        self._lock = Lock()
+
+    def _next_client(self) -> ModelHubClient:
+        with self._lock:
+            return next(self._cycle)
+
+    def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
+        return self._next_client().add_task(payload)
+
+    def search_by_model_id(self, model_id: str) -> dict[str, Any]:
+        return self.clients[0].search_by_model_id(model_id)
+
+    def is_model_processed_for_gpu(self, model_id: str, target_gpu: str) -> bool:
+        return self.clients[0].is_model_processed_for_gpu(model_id, target_gpu)
+
+    def get_verify_result_map(self, model_id: str) -> dict[str, Any]:
+        return self.clients[0].get_verify_result_map(model_id)
+
+    def processed_gpus_for_model(self, model_id: str) -> set[str]:
+        return self.clients[0].processed_gpus_for_model(model_id)
+
+    def list_tasks_page(self, **kwargs: Any) -> dict[str, Any]:
+        return self.clients[0].list_tasks_page(**kwargs)
+
+    def list_tasks(self, **kwargs: Any) -> list[dict[str, Any]]:
+        return self.clients[0].list_tasks(**kwargs)
+
+    def find_recent_task_id(self, model_id: str, gpu_type: str, submitted_after: datetime) -> str | None:
+        return self.clients[0].find_recent_task_id(model_id, gpu_type, submitted_after)
+
+    def count_active_tasks(self, **kwargs: Any) -> int:
+        return self.clients[0].count_active_tasks(**kwargs)
+
+    def active_task_counts(self) -> list[int]:
+        counts: list[int] = []
+        now = datetime.utcnow()
+        begin_time = now - timedelta(days=1)
+        for client in self.clients:
+            counts.append(client.count_active_tasks(begin_time=begin_time, end_time=now, max_count=100))
+        return counts
+
+    def available_submit_slots(self, max_active_per_account: int = 5) -> int:
+        return sum(max(0, max_active_per_account - count) for count in self.active_task_counts())
+
+
 ACTIVE_TASK_STATUSES = {
     "waiting",
     "running",