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 code == 40000 and "正在验证中" 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)