adapt for ModelHub agent platform
This commit is contained in:
CoolBoy
230
modelhub_submmit_api/modelhub_client.py
Normal file
230
modelhub_submmit_api/modelhub_client.py
Normal file
@@ -0,0 +1,230 @@
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user