154 lines
7.2 KiB
Python
154 lines
7.2 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
import logging
|
|
|
|
import requests
|
|
|
|
from app.settings import Settings
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ApiResult:
|
|
result: str
|
|
code: str | None = None
|
|
message: str | None = None
|
|
payload: dict[str, Any] | None = None
|
|
|
|
|
|
class ModelHubClient:
|
|
def __init__(self, settings: Settings):
|
|
self.settings = settings
|
|
self.session = requests.Session()
|
|
|
|
def _headers(self) -> dict[str, str]:
|
|
return {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {self.settings.auth_token}",
|
|
"Xc-Token": self.settings.auth_token,
|
|
}
|
|
|
|
def is_model_absent_from_modelhub(self, model_id: str) -> bool:
|
|
url = self.settings.modelhub_api_base + "/computility/models/list/page/vo"
|
|
data = {"current": 1, "pageSize": 20, "searchText": model_id}
|
|
try:
|
|
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
|
|
response.raise_for_status()
|
|
total = int(response.json().get("data", {}).get("total", 0))
|
|
return total == 0
|
|
except Exception as exc:
|
|
LOG.warning("modelhub existence check failed for %s: %s", model_id, exc)
|
|
return True
|
|
|
|
def list_modelhub_models(self, current: int = 1, page_size: int = 100, search_text: str = "") -> dict[str, Any]:
|
|
url = self.settings.modelhub_api_base + "/computility/models/list/page/vo"
|
|
data = {"current": current, "pageSize": page_size, "searchText": search_text}
|
|
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
|
|
response.raise_for_status()
|
|
body = response.json()
|
|
if body.get("code") != 0:
|
|
raise RuntimeError(body.get("message") or body.get("msg") or "list_modelhub_models failed")
|
|
return body.get("data") or {}
|
|
|
|
def list_success_download_tasks(self, current: int = 1, page_size: int = 50) -> dict[str, Any]:
|
|
url = self.settings.modelhub_adminapi_base + "/async/task/model-download-task"
|
|
params = {"current": current, "pageSize": page_size, "status": "SUCCESS"}
|
|
response = self.session.get(url, headers=self._headers(), params=params, timeout=self.settings.request_timeout_seconds)
|
|
response.raise_for_status()
|
|
body = response.json()
|
|
if body.get("code") != 0:
|
|
raise RuntimeError(body.get("message") or "list_success_download_tasks failed")
|
|
return body.get("data") or {}
|
|
|
|
def sync_from_huggingface(self, model_id: str) -> ApiResult:
|
|
url = self.settings.modelhub_adminapi_base + "/computility/models/sync-from-hugging-face"
|
|
data = {"modelId": model_id, "forceUpdate": False, "operatorEmail": self.settings.email}
|
|
try:
|
|
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
|
|
payload = response.json()
|
|
return ApiResult("success" if payload.get("code") == 0 else "failed", str(payload.get("code")), payload.get("message"), payload)
|
|
except Exception as exc:
|
|
return ApiResult("failed", None, f"sync exception: {exc}", None)
|
|
|
|
def create_download_task(self, model_id: str, source: str = "HUGGING_FACE", weight_file: str | None = None) -> ApiResult:
|
|
url = self.settings.modelhub_url + "/adminApi/async/task/model-download-task"
|
|
data: dict[str, Any] = {
|
|
"hfToken": self.settings.hf_token,
|
|
"modelId": model_id,
|
|
"source": source,
|
|
"stillDownloadAlreadySuccessDownloadedModel": False,
|
|
}
|
|
if weight_file:
|
|
data["allowPatterns"] = [weight_file]
|
|
try:
|
|
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
|
|
payload = response.json()
|
|
code = payload.get("code")
|
|
msg = payload.get("message", "")
|
|
if code == 0:
|
|
result = "created"
|
|
elif code == 40000:
|
|
result = "running_elsewhere"
|
|
elif code == 60004:
|
|
result = "already_downloaded"
|
|
elif code == 60005:
|
|
result = "download_pool_full"
|
|
else:
|
|
result = "failed"
|
|
return ApiResult(result, str(code), msg, payload)
|
|
except Exception as exc:
|
|
return ApiResult("failed", None, f"download exception: {exc}", None)
|
|
|
|
def get_download_status(self, model_id: str) -> str:
|
|
url = self.settings.modelhub_url + "/adminApi/async/task/model-download-task"
|
|
params = {"modelId": model_id, "userEmail": self.settings.email, "userId": self.settings.user_id}
|
|
try:
|
|
response = self.session.get(url, headers=self._headers(), params=params, timeout=self.settings.request_timeout_seconds)
|
|
response.raise_for_status()
|
|
records = response.json().get("data", {}).get("records", [])
|
|
if records:
|
|
return records[0].get("status") or "None"
|
|
return "None"
|
|
except Exception as exc:
|
|
LOG.warning("download status check failed for %s: %s", model_id, exc)
|
|
return "QUERY_FAILED"
|
|
|
|
def build_contest_payload(self, model_id: str, gpu_type: str, config: str) -> dict[str, Any]:
|
|
payload: dict[str, Any] = {
|
|
"contestApiToken": self.settings.contest_api_token,
|
|
"contributors": self.settings.contributors,
|
|
"gpuTypes": [gpu_type],
|
|
"modelId": model_id,
|
|
"submissionConfig": [{"gpuType": gpu_type, "config": config}],
|
|
"taskType": "text-generation",
|
|
}
|
|
if self.settings.inject_strategy_id:
|
|
payload[self.settings.contest_task_strategy_field] = self.settings.strategy_id
|
|
return payload
|
|
|
|
def create_contest_task(self, model_id: str, gpu_type: str, config: str) -> ApiResult:
|
|
payload = self.build_contest_payload(model_id, gpu_type, config)
|
|
if self.settings.submit_dry_run:
|
|
LOG.info("dry-run submit payload: model=%s gpu=%s strategy_field=%s", model_id, gpu_type, self.settings.contest_task_strategy_field)
|
|
return ApiResult("success", "DRY_RUN", "dry run", {"payload": payload})
|
|
url = self.settings.modelhub_adminapi_base + "/async/task/create-contest-task"
|
|
try:
|
|
response = self.session.post(url, headers=self._headers(), json=payload, timeout=self.settings.request_timeout_seconds)
|
|
body = response.json()
|
|
code = body.get("code")
|
|
message = body.get("message", "")
|
|
if code == 0 and message == "ok":
|
|
result = "success"
|
|
elif "请勿重复提交" in message or "已经被验证成功" in message or "正在验证中" in message:
|
|
result = "conflict"
|
|
elif code == 60007:
|
|
result = "queue_full"
|
|
else:
|
|
result = "failed"
|
|
return ApiResult(result, str(code), message, body)
|
|
except Exception as exc:
|
|
return ApiResult("failed", None, f"submit exception: {exc}", None)
|