submmit/modelhub_submmit_api/modelhub_client.py

from __future__ import annotations

import hashlib
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timedelta
from typing import Any

from common import format_modelhub_datetime, parse_datetime, runtime_instance_id
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 EMBEDDED_MODELHUB_XC_TOKEN
        if not self.token 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


CAPACITY_ERROR_MARKERS = (
    "达到上限",
    "达上限",
    "任务数量已达",
    "队列已满",
    "queue is full",
    "queue full",
    "capacity",
    "too many active",
    "active task limit",
)


def is_capacity_error(error: ModelHubAPIError) -> bool:
    if error.code in {409, 429}:
        return True
    message = str(error).strip().lower()
    if isinstance(error.payload, dict):
        message = f"{message} {error.payload.get('message') or ''}".lower()
    return any(marker in message for marker in CAPACITY_ERROR_MARKERS)


class ModelHubClientPool:
    def __init__(
        self,
        clients: list[ModelHubClient],
        *,
        active_task_cap: int | None = None,
        active_counts_ttl: float | None = None,
        reservation_ttl: float | None = None,
        instance_id: str | None = None,
    ) -> None:
        if not clients:
            raise ValueError("At least one ModelHub client is required")
        self.clients = clients
        configured_cap = active_task_cap if active_task_cap is not None else os.getenv("MODELHUB_AGENT_ACTIVE_TASK_CAP", "100")
        self.active_task_cap = max(1, int(configured_cap))
        configured_ttl = active_counts_ttl if active_counts_ttl is not None else os.getenv("MODELHUB_AGENT_ACTIVE_COUNTS_TTL_SECONDS", "15")
        configured_reservation_ttl = (
            reservation_ttl
            if reservation_ttl is not None
            else os.getenv("MODELHUB_AGENT_RESERVATION_TTL_SECONDS", "120")
        )
        self._active_counts_ttl = max(1.0, float(configured_ttl))
        self._reservation_ttl = max(self._active_counts_ttl * 2, float(configured_reservation_ttl))
        self._remote_counts: list[int] = [0 for _ in clients]
        self._active_refresh_at: float = 0.0
        self._counts_initialized = False
        self._state_lock = threading.Lock()
        self._refresh_lock = threading.Lock()
        self._reservations: list[dict[int, dict[str, Any]]] = [{} for _ in clients]
        self._reservation_sequence = 0
        identity = instance_id or runtime_instance_id()
        identity_hash = hashlib.sha256(identity.encode("utf-8")).hexdigest()
        self._selection_cursor = int(identity_hash[:12], 16) % len(clients)
        self._verify_cache: dict[str, tuple[float, dict[str, Any]]] = {}
        self._verify_cache_ttl = max(1.0, float(os.getenv("MODELHUB_AGENT_VERIFY_CACHE_TTL_SECONDS", "30")))
        # Single reader client to avoid fanout on read operations
        self._reader = clients[0]

    def _safe_count_active_tasks(self, client: ModelHubClient) -> int:
        try:
            return client.count_active_tasks(max_count=self.active_task_cap, page_size=200)
        except Exception:
            return self.active_task_cap

    def _safe_search_by_model_id(self, client: ModelHubClient, model_id: str) -> dict[str, Any]:
        try:
            payload = client.search_by_model_id(model_id)
            return payload if isinstance(payload, dict) else {}
        except Exception:
            return {}

    def _safe_list_tasks(self, client: ModelHubClient, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
        try:
            return client.list_tasks(**kwargs)
        except Exception:
            return []

    def _counts_are_fresh_locked(self, now: float) -> bool:
        return self._counts_initialized and (now - self._active_refresh_at) < self._active_counts_ttl

    def _refresh_active_counts(self, *, force: bool = False) -> None:
        now = time.monotonic()
        with self._state_lock:
            if not force and self._counts_are_fresh_locked(now):
                return

        # Do not hold the scheduler lock during remote I/O. One refresher is
        # enough; all submitting threads can continue using their reservations.
        with self._refresh_lock:
            now = time.monotonic()
            with self._state_lock:
                if not force and self._counts_are_fresh_locked(now):
                    return

            def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int]:
                return index, self._safe_count_active_tasks(client)

            results: list[tuple[int, int]] = []
            with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
                futures = {
                    executor.submit(_to_indexed_result, index, client): index
                    for index, client in enumerate(self.clients)
                }
                for future in as_completed(futures):
                    index = futures[future]
                    try:
                        results.append(future.result())
                    except Exception:
                        results.append((index, self.active_task_cap))

            refreshed_at = time.monotonic()
            with self._state_lock:
                for index, count in results:
                    old_remote_count = self._remote_counts[index]
                    new_remote_count = min(self.active_task_cap, max(0, int(count)))
                    acknowledged = max(0, new_remote_count - old_remote_count)
                    completed_reservations = sorted(
                        (
                            (reservation_id, reservation)
                            for reservation_id, reservation in self._reservations[index].items()
                            if not reservation["inflight"]
                        ),
                        key=lambda item: item[1]["updated_at"],
                    )
                    for reservation_id, _reservation in completed_reservations[:acknowledged]:
                        self._reservations[index].pop(reservation_id, None)

                    # A remote count can stay flat when one old task finishes as
                    # one new task appears. Expiry prevents that net-zero update
                    # from reserving a slot forever.
                    for reservation_id, reservation in list(self._reservations[index].items()):
                        if reservation["inflight"]:
                            continue
                        if refreshed_at - reservation["updated_at"] >= self._reservation_ttl:
                            self._reservations[index].pop(reservation_id, None)
                    self._remote_counts[index] = new_remote_count
                self._counts_initialized = True
                self._active_refresh_at = refreshed_at

    def _effective_count_locked(self, index: int) -> int:
        return self._remote_counts[index] + len(self._reservations[index])

    def _counts_snapshot_locked(self) -> list[int]:
        return [min(self.active_task_cap, self._effective_count_locked(index)) for index in range(len(self.clients))]

    def active_task_counts(self) -> list[int]:
        self._refresh_active_counts()
        with self._state_lock:
            return self._counts_snapshot_locked()

    def available_submit_slots(self) -> int:
        self._refresh_active_counts()
        with self._state_lock:
            return sum(
                max(0, self.active_task_cap - self._effective_count_locked(index))
                for index in range(len(self.clients))
            )

    def list_tasks(self, **kwargs):  # noqa: ANN003, ANN001
        # Default: use only the reader client to avoid fanout amplification.
        # The fanout_all parameter allows explicit cross-account merging when needed.
        fanout_all = kwargs.pop("_fanout_all", False)
        if not fanout_all:
            return self._safe_list_tasks(self._reader, kwargs)

        merged: list[dict[str, Any]] = []
        seen: set[str] = set()
        with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
            futures = {executor.submit(self._safe_list_tasks, client, kwargs): index for index, client in enumerate(self.clients)}
            for future in as_completed(futures):
                try:
                    tasks = future.result()
                except Exception:
                    tasks = []
                for task in tasks:
                    task_id = str(task.get("taskId")) if task.get("taskId") is not None else None
                    if task_id and task_id in seen:
                        continue
                    if task_id:
                        seen.add(task_id)
                    merged.append(task)
        return merged

    def _read_from_cache(self, model_id: str) -> dict[str, Any] | None:
        cached = self._verify_cache.get(model_id)
        if cached is None:
            return None
        cached_at, payload = cached
        if time.monotonic() - cached_at >= self._verify_cache_ttl:
            self._verify_cache.pop(model_id, None)
            return None
        return dict(payload)

    def _write_to_cache(self, model_id: str, payload: dict[str, Any]) -> None:
        self._verify_cache[model_id] = (time.monotonic(), payload)

    def begin_cycle(self) -> None:
        """Drop model verification cache entries from the previous scan cycle."""
        with self._state_lock:
            self._verify_cache.clear()

    def search_by_model_id(self, model_id: str) -> dict[str, Any]:
        # Check cache first (per-cycle cache to avoid repeated API calls for the same model)
        with self._state_lock:
            cached = self._read_from_cache(model_id)
        if cached is not None:
            return cached

        # Only query ONE client (the reader) instead of fanning out to all clients.
        # verifyResult is model-specific platform data, not account-specific.
        response = self._safe_search_by_model_id(self._reader, model_id)
        if not isinstance(response, dict):
            response = {"code": 0, "data": {"verifyResult": {}}}

        with self._state_lock:
            self._write_to_cache(model_id, response)
        return response

    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 _reserve_account(self, excluded: set[int]) -> tuple[int, int] | None:
        with self._state_lock:
            remaining_by_index = {
                index: self.active_task_cap - self._effective_count_locked(index)
                for index in range(len(self.clients))
                if index not in excluded
            }
            usable = {index: remaining for index, remaining in remaining_by_index.items() if remaining > 0}
            if not usable:
                return None
            best_remaining = max(usable.values())
            tied = [index for index, remaining in usable.items() if remaining == best_remaining]
            selected_index = min(
                tied,
                key=lambda index: (index - self._selection_cursor) % len(self.clients),
            )
            self._selection_cursor = (selected_index + 1) % len(self.clients)
            self._reservation_sequence += 1
            reservation_id = self._reservation_sequence
            self._reservations[selected_index][reservation_id] = {
                "inflight": True,
                "updated_at": time.monotonic(),
            }
            return selected_index, reservation_id

    def _finish_reservation(self, index: int, reservation_id: int, *, succeeded: bool) -> None:
        with self._state_lock:
            reservation = self._reservations[index].get(reservation_id)
            if reservation is None:
                return
            if not succeeded:
                self._reservations[index].pop(reservation_id, None)
                return
            reservation["inflight"] = False
            reservation["updated_at"] = time.monotonic()

    def _mark_account_saturated(self, index: int) -> None:
        with self._state_lock:
            self._remote_counts[index] = self.active_task_cap
            self._counts_initialized = True
            self._active_refresh_at = time.monotonic()

    def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
        self._refresh_active_counts()
        attempted_accounts: set[int] = set()
        forced_refresh_done = False
        last_capacity_error: ModelHubAPIError | None = None

        while True:
            reservation = self._reserve_account(attempted_accounts)
            if reservation is None:
                if not forced_refresh_done:
                    self._refresh_active_counts(force=True)
                    forced_refresh_done = True
                    continue
                if last_capacity_error is not None:
                    raise last_capacity_error
                raise ModelHubAPIError(
                    f"当前等待中或运行中的异步模型验证任务数量已达上限（{self.active_task_cap}）"
                )

            selected_index, reservation_id = reservation
            selected_client = self.clients[selected_index]
            try:
                response = selected_client.add_task(payload)
            except ModelHubAPIError as exc:
                self._finish_reservation(selected_index, reservation_id, succeeded=False)
                if not is_capacity_error(exc):
                    raise
                # Another process may have filled this account after our count
                # refresh. Mark it full locally and immediately try another one.
                self._mark_account_saturated(selected_index)
                attempted_accounts.add(selected_index)
                last_capacity_error = exc
                continue
            except Exception:
                self._finish_reservation(selected_index, reservation_id, succeeded=False)
                raise

            self._finish_reservation(selected_index, reservation_id, succeeded=True)
            return response

    def list_tasks_page(self, **kwargs):  # noqa: ANN003, ANN001
        """Single-page task listing via the reader client (no fanout)."""
        return self._reader.list_tasks_page(**kwargs)

    def find_recent_task_id(self, model_id: str, gpu_type: str, submitted_after: datetime) -> str | None:
        for client in self.clients:
            task_id = client.find_recent_task_id(model_id, gpu_type, submitted_after)
            if task_id is not None:
                return task_id
        return None