Initial vLLM agent strategy

app/clients/modelhub.py

@@ -0,0 +1,153 @@
+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)