231 lines
7.6 KiB
Python
231 lines
7.6 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import os
|
||
|
|
from datetime import datetime, timedelta
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from common import format_modelhub_datetime, parse_datetime
|
||
|
|
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")
|
||
|
|
if not self.token and not os.getenv("STRATEGY_ID") 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,
|
||
|
|
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
|
||
|
|
|
||
|
|
|
||
|
|
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
|
||
|
|
|