submmit/modelhub_submmit_api/modelhub_client.py

from __future__ import annotations

import os
from datetime import datetime, timedelta
from itertools import cycle
from threading import Lock
from typing import Any

from common import format_modelhub_datetime, parse_datetime
from defaults import EMBEDDED_MODELHUB_XC_TOKEN
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")
            or os.getenv("MODELHUB_TOKEN")
            or EMBEDDED_MODELHUB_XC_TOKEN
        )
        self.jwt_token = os.getenv("MODELHUB_JWT_TOKEN") or os.getenv("JWT_TOKEN")
        if not self.token and not self.jwt_token and http_client is None:
            raise ValueError("ModelHub token is required. Set MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN.")
        default_headers = {}
        if self.token:
            default_headers["Xc-Token"] = self.token
        if self.jwt_token:
            default_headers["Authorization"] = f"Bearer {self.jwt_token}"
        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


class MultiModelHubClient:
    def __init__(
        self,
        *,
        tokens: list[str],
        base_url: str = "https://modelhub.org.cn",
        timeout: int = 30,
        retries: int = 2,
    ) -> None:
        deduped: list[str] = []
        for token in tokens:
            clean = token.strip()
            if clean and clean not in deduped:
                deduped.append(clean)
        if not deduped:
            raise ValueError("At least one ModelHub token is required")

        self.clients = [
            ModelHubClient(token=token, base_url=base_url, timeout=timeout, retries=retries)
            for token in deduped
        ]
        self._cycle = cycle(self.clients)
        self._lock = Lock()

    def _next_client(self) -> ModelHubClient:
        with self._lock:
            return next(self._cycle)

    def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
        return self._next_client().add_task(payload)

    def search_by_model_id(self, model_id: str) -> dict[str, Any]:
        return self.clients[0].search_by_model_id(model_id)

    def is_model_processed_for_gpu(self, model_id: str, target_gpu: str) -> bool:
        return self.clients[0].is_model_processed_for_gpu(model_id, target_gpu)

    def get_verify_result_map(self, model_id: str) -> dict[str, Any]:
        return self.clients[0].get_verify_result_map(model_id)

    def processed_gpus_for_model(self, model_id: str) -> set[str]:
        return self.clients[0].processed_gpus_for_model(model_id)

    def list_tasks_page(self, **kwargs: Any) -> dict[str, Any]:
        return self.clients[0].list_tasks_page(**kwargs)

    def list_tasks(self, **kwargs: Any) -> list[dict[str, Any]]:
        return self.clients[0].list_tasks(**kwargs)

    def find_recent_task_id(self, model_id: str, gpu_type: str, submitted_after: datetime) -> str | None:
        return self.clients[0].find_recent_task_id(model_id, gpu_type, submitted_after)

    def count_active_tasks(self, **kwargs: Any) -> int:
        return self.clients[0].count_active_tasks(**kwargs)

    def active_task_counts(self) -> list[int]:
        counts: list[int] = []
        now = datetime.utcnow()
        begin_time = now - timedelta(days=1)
        for client in self.clients:
            counts.append(client.count_active_tasks(begin_time=begin_time, end_time=now, max_count=100))
        return counts

    def available_submit_slots(self, max_active_per_account: int = 5) -> int:
        return sum(max(0, max_active_per_account - count) for count in self.active_task_counts())


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