311 lines
11 KiB
Python
311 lines
11 KiB
Python
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
|
|
from defaults import EMBEDDED_MODELHUB_XC_TOKEN
|
|
from http_json import HttpJsonError, JsonHttpClient
|
|
|
|
|
|
class ModelHubAPIError(RuntimeError):
|
|
def __init__(self, message: str, *, code: int | None = None, payload: Any = None) -> None:
|
|
super().__init__(message)
|
|
self.code = code
|
|
self.payload = payload
|
|
|
|
|
|
class ModelHubClient:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
token: str | None = None,
|
|
base_url: str = "https://modelhub.org.cn",
|
|
timeout: int = 30,
|
|
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.http_client = http_client or JsonHttpClient(
|
|
base_url=base_url,
|
|
default_headers=default_headers,
|
|
timeout=timeout,
|
|
retries=retries,
|
|
)
|
|
|
|
def search_by_model_id(self, model_id: str) -> dict[str, Any]:
|
|
return self._request("GET", "/api/computility/models/search-by-model-id", query={"modelId": model_id})
|
|
|
|
def is_model_processed_for_gpu(self, model_id: str, target_gpu: str) -> bool:
|
|
gpu_result = self.get_verify_result_map(model_id).get(target_gpu)
|
|
if gpu_result is None:
|
|
return False
|
|
return "result" in gpu_result or "records" in gpu_result
|
|
|
|
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 list_tasks_page(
|
|
self,
|
|
*,
|
|
current: int = 1,
|
|
page_size: int = 50,
|
|
only_mine: bool = True,
|
|
begin_time: datetime | None = None,
|
|
end_time: datetime | None = None,
|
|
gpu_type: str | None = None,
|
|
model_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
return self._request(
|
|
"GET",
|
|
"/api/adapt/task/page",
|
|
query={
|
|
"current": current,
|
|
"pageSize": page_size,
|
|
"onlyMine": str(only_mine).lower(),
|
|
"beginTime": format_modelhub_datetime(begin_time) if begin_time else None,
|
|
"endTime": format_modelhub_datetime(end_time) if end_time else None,
|
|
"gpuType": gpu_type,
|
|
"modelId": model_id,
|
|
},
|
|
)
|
|
|
|
def list_tasks(
|
|
self,
|
|
*,
|
|
page_size: int = 50,
|
|
only_mine: bool = True,
|
|
begin_time: datetime | None = None,
|
|
end_time: datetime | None = None,
|
|
gpu_type: str | None = None,
|
|
model_id: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
current = 1
|
|
records: list[dict[str, Any]] = []
|
|
while True:
|
|
page = self.list_tasks_page(
|
|
current=current,
|
|
page_size=page_size,
|
|
only_mine=only_mine,
|
|
begin_time=begin_time,
|
|
end_time=end_time,
|
|
gpu_type=gpu_type,
|
|
model_id=model_id,
|
|
)
|
|
page_data = page.get("data") or {}
|
|
page_records = page_data.get("records") or []
|
|
records.extend(page_records)
|
|
pages = int(page_data.get("pages") or 0)
|
|
if pages <= current or not page_records:
|
|
break
|
|
current += 1
|
|
return records
|
|
|
|
def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
return self._request("POST", "/api/adapt/task/add", data=payload)
|
|
|
|
def find_recent_task_id(self, model_id: str, gpu_type: str, submitted_after: datetime) -> str | None:
|
|
recent_tasks = self.list_tasks(
|
|
page_size=20,
|
|
only_mine=True,
|
|
begin_time=submitted_after - timedelta(hours=1),
|
|
end_time=submitted_after + timedelta(hours=6),
|
|
gpu_type=gpu_type,
|
|
model_id=model_id,
|
|
)
|
|
if not recent_tasks:
|
|
return None
|
|
recent_tasks.sort(key=lambda task: parse_datetime(task.get("updateTime")) or submitted_after, reverse=True)
|
|
return str(recent_tasks[0].get("taskId")) if recent_tasks[0].get("taskId") is not None else None
|
|
|
|
def count_active_tasks(
|
|
self,
|
|
*,
|
|
page_size: int = 100,
|
|
only_mine: bool = True,
|
|
begin_time: datetime | None = None,
|
|
end_time: datetime | None = None,
|
|
max_count: int | None = None,
|
|
) -> int:
|
|
"""Count active tasks with early pagination termination.
|
|
|
|
Unlike list_tasks() which fetches all pages, this method stops
|
|
fetching pages as soon as max_count active tasks are found.
|
|
"""
|
|
max_count_int: int | None = max_count
|
|
if max_count_int is not None:
|
|
max_count_int = max(0, int(max_count_int))
|
|
if max_count_int == 0:
|
|
return 0
|
|
|
|
active_count = 0
|
|
current = 1
|
|
while True:
|
|
page = self.list_tasks_page(
|
|
current=current,
|
|
page_size=page_size,
|
|
only_mine=only_mine,
|
|
begin_time=begin_time,
|
|
end_time=end_time,
|
|
)
|
|
page_data = page.get("data") or {}
|
|
page_records = page_data.get("records") or []
|
|
|
|
for task in page_records:
|
|
if is_active_task(task):
|
|
active_count += 1
|
|
if max_count_int is not None and active_count >= max_count_int:
|
|
return active_count
|
|
|
|
pages = int(page_data.get("pages") or 0)
|
|
if pages <= current or not page_records:
|
|
break
|
|
current += 1
|
|
|
|
return active_count
|
|
|
|
def _request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
*,
|
|
query: dict[str, Any] | None = None,
|
|
data: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
try:
|
|
payload = self.http_client.request_json(method, path, query=query, data=data)
|
|
except HttpJsonError as exc:
|
|
if isinstance(exc.payload, dict) and "message" in exc.payload:
|
|
raise ModelHubAPIError(exc.payload["message"], payload=exc.payload, code=exc.status_code) from exc
|
|
raise ModelHubAPIError(str(exc), payload=exc.payload, code=exc.status_code) from exc
|
|
if not isinstance(payload, dict):
|
|
raise ModelHubAPIError("Unexpected API response shape", payload=payload)
|
|
code = payload.get("code")
|
|
if code != 0:
|
|
raise ModelHubAPIError(payload.get("message") or "ModelHub API request failed", code=code, payload=payload)
|
|
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",
|
|
"processing",
|
|
"queued",
|
|
"pending",
|
|
"validating",
|
|
"submitting",
|
|
"initializing",
|
|
}
|
|
|
|
TERMINAL_TASK_STATUSES = {
|
|
"success",
|
|
"failed",
|
|
"error",
|
|
"cancelled",
|
|
"canceled",
|
|
"rejected",
|
|
"timeout",
|
|
"completed",
|
|
"complete",
|
|
"done",
|
|
}
|
|
|
|
|
|
def is_active_task(task: dict[str, Any]) -> bool:
|
|
status = str(task.get("status") or "").strip().lower()
|
|
if not status:
|
|
return False
|
|
if status in ACTIVE_TASK_STATUSES:
|
|
return True
|
|
if status in TERMINAL_TASK_STATUSES:
|
|
return False
|
|
return True
|