Rebuild agent on original pooled runner
This commit is contained in:
CoolBoy
@@ -1,9 +1,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
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
|
||||
@@ -28,21 +29,10 @@ class ModelHubClient:
|
||||
retries: int = 2,
|
||||
http_client: JsonHttpClient | None = None,
|
||||
) -> None:
|
||||
self.token = (
|
||||
token
|
||||
or os.getenv("MODELHUB_XC_TOKEN")
|
||||
or os.getenv("XC_TOKEN")
|
||||
or os.getenv("MODELHUB_TOKEN")
|
||||
or EMBEDDED_MODELHUB_XC_TOKEN
|
||||
)
|
||||
self.jwt_token = os.getenv("MODELHUB_JWT_TOKEN") or os.getenv("JWT_TOKEN")
|
||||
if not self.token and not self.jwt_token and http_client is None:
|
||||
raise ValueError("ModelHub token is required. Set MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN.")
|
||||
default_headers = {}
|
||||
if self.token:
|
||||
default_headers["Xc-Token"] = self.token
|
||||
if self.jwt_token:
|
||||
default_headers["Authorization"] = f"Bearer {self.jwt_token}"
|
||||
self.token = token or os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN") or EMBEDDED_MODELHUB_XC_TOKEN
|
||||
if not self.token and http_client is None:
|
||||
raise ValueError("ModelHub token is required. Set MODELHUB_XC_TOKEN or XC_TOKEN.")
|
||||
default_headers = {"Xc-Token": self.token} if self.token else {}
|
||||
self.http_client = http_client or JsonHttpClient(
|
||||
base_url=base_url,
|
||||
default_headers=default_headers,
|
||||
@@ -207,73 +197,6 @@ 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",
|
||||
@@ -308,3 +231,175 @@ def is_active_task(task: dict[str, Any]) -> bool:
|
||||
if status in TERMINAL_TASK_STATUSES:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class ModelHubClientPool:
|
||||
def __init__(self, clients: list[ModelHubClient], *, active_task_cap: int = 100) -> 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]
|
||||
self._active_refresh_at: float = 0.0
|
||||
self._active_counts_ttl: float = 30.0
|
||||
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]] = {}
|
||||
# Single reader client to avoid fanout on read operations
|
||||
self._reader = clients[0]
|
||||
|
||||
def _safe_count_active_tasks(self, client: ModelHubClient) -> int:
|
||||
try:
|
||||
return client.count_active_tasks(max_count=self.active_task_cap, page_size=200)
|
||||
except Exception:
|
||||
return self.active_task_cap
|
||||
|
||||
def _safe_search_by_model_id(self, client: ModelHubClient, model_id: str) -> dict[str, Any]:
|
||||
try:
|
||||
payload = client.search_by_model_id(model_id)
|
||||
return payload if isinstance(payload, dict) else {}
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
def _safe_list_tasks(self, client: ModelHubClient, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
try:
|
||||
return client.list_tasks(**kwargs)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
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
|
||||
|
||||
def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int]:
|
||||
return index, self._safe_count_active_tasks(client)
|
||||
|
||||
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))
|
||||
|
||||
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()
|
||||
|
||||
def active_task_counts(self) -> list[int]:
|
||||
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)
|
||||
|
||||
def available_submit_slots(self) -> int:
|
||||
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)
|
||||
|
||||
def list_tasks(self, **kwargs): # noqa: ANN003, ANN001
|
||||
# Default: use only the reader client to avoid fanout amplification.
|
||||
# The fanout_all parameter allows explicit cross-account merging when needed.
|
||||
fanout_all = kwargs.pop("_fanout_all", False)
|
||||
if not fanout_all:
|
||||
return self._safe_list_tasks(self._reader, kwargs)
|
||||
|
||||
merged: list[dict[str, Any]] = []
|
||||
seen: set[str] = set()
|
||||
with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
|
||||
futures = {executor.submit(self._safe_list_tasks, client, kwargs): index for index, client in enumerate(self.clients)}
|
||||
for future in as_completed(futures):
|
||||
try:
|
||||
tasks = future.result()
|
||||
except Exception:
|
||||
tasks = []
|
||||
for task in tasks:
|
||||
task_id = str(task.get("taskId")) if task.get("taskId") is not None else None
|
||||
if task_id and task_id in seen:
|
||||
continue
|
||||
if task_id:
|
||||
seen.add(task_id)
|
||||
merged.append(task)
|
||||
return merged
|
||||
|
||||
def _read_from_cache(self, model_id: str) -> dict[str, Any] | None:
|
||||
cached = self._verify_cache.get(model_id)
|
||||
if cached is None:
|
||||
return None
|
||||
return dict(cached)
|
||||
|
||||
def _write_to_cache(self, model_id: str, payload: dict[str, Any]) -> None:
|
||||
self._verify_cache[model_id] = payload
|
||||
|
||||
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)
|
||||
with self._state_lock:
|
||||
cached = self._read_from_cache(model_id)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
# Only query ONE client (the reader) instead of fanning out to all clients.
|
||||
# verifyResult is model-specific platform data, not account-specific.
|
||||
response = self._safe_search_by_model_id(self._reader, model_id)
|
||||
if not isinstance(response, dict):
|
||||
response = {"code": 0, "data": {"verifyResult": {}}}
|
||||
|
||||
with self._state_lock:
|
||||
self._write_to_cache(model_id, response)
|
||||
return response
|
||||
|
||||
def get_verify_result_map(self, model_id: str) -> dict[str, Any]:
|
||||
payload = self.search_by_model_id(model_id)
|
||||
return ((payload.get("data") or {}).get("verifyResult") or {})
|
||||
|
||||
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]:
|
||||
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:
|
||||
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
|
||||
|
||||
def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN001
|
||||
"""Single-page task listing via the reader client (no fanout)."""
|
||||
return self._reader.list_tasks_page(**kwargs)
|
||||
|
||||
def find_recent_task_id(self, model_id: str, gpu_type: str, submitted_after: datetime) -> str | None:
|
||||
for client in self.clients:
|
||||
task_id = client.find_recent_task_id(model_id, gpu_type, submitted_after)
|
||||
if task_id is not None:
|
||||
return task_id
|
||||
return None
|
||||
|
||||
Reference in New Issue
Block a user