submmit/modelhub_submmit_api/official_capabilities.py

from __future__ import annotations

from datetime import datetime, timedelta
from pathlib import Path
import threading
from typing import Any

from common import parse_datetime, read_json, utc_now, write_json


OFFICIAL_CAPABILITY_VERSION = 1
DEFAULT_OFFICIAL_CAPABILITIES_PATH = Path(".modelhub_state/official_capabilities.json")
DEFAULT_MODEL_GPU_CACHE_LIMIT = 750


class OfficialCapabilityUnavailable(RuntimeError):
    pass


def _fresh(value: Any, now: datetime, seconds: int) -> bool:
    parsed = parse_datetime(value)
    return parsed is not None and now - parsed <= timedelta(seconds=max(1, seconds))


def _extract_task_names(payload: Any) -> list[str]:
    data = payload.get("data") if isinstance(payload, dict) else payload
    result: list[str] = []

    def visit(value: Any) -> None:
        if isinstance(value, str):
            if value.strip() and value.strip() not in result:
                result.append(value.strip())
            return
        if isinstance(value, list):
            for item in value:
                visit(item)
            return
        if not isinstance(value, dict):
            return
        for key in ("taskType", "taskTypeName", "value", "code"):
            candidate = value.get(key)
            if isinstance(candidate, str) and candidate.strip() and candidate.strip() not in result:
                result.append(candidate.strip())
        children = value.get("children")
        if not children:
            candidate = value.get("name")
            if isinstance(candidate, str) and candidate.strip() and candidate.strip() not in result:
                result.append(candidate.strip())
        for key in ("children", "items", "records", "list"):
            visit(value.get(key))

    visit(data)
    return result


class OfficialCapabilityRegistry:
    def __init__(self, path: Path | str = DEFAULT_OFFICIAL_CAPABILITIES_PATH, *, log_fn=None) -> None:
        self.path = Path(path)
        self.log = log_fn or (lambda message: print(message, flush=True))
        self.state: dict[str, Any] = {}
        self.ready = False
        self.pause_reason: str | None = None
        self._lock = threading.Lock()

    @staticmethod
    def _prune_model_gpu_cache(state: dict[str, Any]) -> None:
        cache = state.get("modelGpuTaskTypes")
        if not isinstance(cache, dict) or len(cache) <= DEFAULT_MODEL_GPU_CACHE_LIMIT:
            return
        ordered = sorted(
            cache.items(),
            key=lambda pair: parse_datetime((pair[1] or {}).get("updatedAt")) or datetime.min.replace(tzinfo=utc_now().tzinfo),
            reverse=True,
        )
        state["modelGpuTaskTypes"] = dict(ordered[:DEFAULT_MODEL_GPU_CACHE_LIMIT])

    def _load(self) -> dict[str, Any]:
        try:
            value = read_json(self.path)
        except (FileNotFoundError, ValueError):
            return {}
        if not isinstance(value, dict) or int(value.get("version") or 0) != OFFICIAL_CAPABILITY_VERSION:
            return {}
        return value

    def prepare(self, client: Any, *, fallback_gpus: list[str], task_types: list[str], now: datetime | None = None) -> dict[str, Any]:
        now = now or utc_now()
        state = self._load()
        self._prune_model_gpu_cache(state)
        catalog_usable = _fresh(state.get("catalogUpdatedAt"), now, 1800)
        task_tree_usable = _fresh(state.get("taskTreeUpdatedAt"), now, 7 * 86400)
        errors: list[str] = []
        try:
            machines = client.list_machine_info()
            catalog: dict[str, dict[str, Any]] = {}
            for item in machines:
                gpu = str(item.get("gpuType") or item.get("gpuTypeName") or "").strip()
                if gpu:
                    catalog[gpu] = {
                        "gpu": gpu,
                        "canVerify": item.get("canVerify"),
                        "maxConcurrentTasks": item.get("maxConcurrentTasks"),
                    }
            if not catalog:
                raise OfficialCapabilityUnavailable("official GPU catalog is empty")
            state["gpuCatalog"] = catalog
            state["catalogUpdatedAt"] = now.isoformat()
            catalog_usable = True
        except Exception as exc:
            errors.append(f"catalog:{type(exc).__name__}:{exc}")

        try:
            payload = client.list_task_levels()
            dynamic_tasks = _extract_task_names(payload)
            if not dynamic_tasks:
                raise OfficialCapabilityUnavailable("official task tree is empty")
            state["taskTree"] = payload.get("data") if isinstance(payload, dict) else payload
            state["discoveredTaskTypes"] = dynamic_tasks
            state["taskTreeUpdatedAt"] = now.isoformat()
            task_tree_usable = True
        except Exception as exc:
            errors.append(f"task_tree:{type(exc).__name__}:{exc}")

        state.update(
            {
                "version": OFFICIAL_CAPABILITY_VERSION,
                "generatedAt": now.isoformat(),
                "configuredTaskTypes": list(task_types),
                "errors": errors,
            }
        )
        self.state = state
        self.ready = bool(catalog_usable and task_tree_usable)
        self.pause_reason = None if self.ready else "critical_official_signal_unavailable"
        write_json(self.path, state)
        eligible = self.eligible_gpus(fallback_gpus)
        self.log(
            f"[official] ready={str(self.ready).lower()} catalog={len(state.get('gpuCatalog') or {})} "
            f"task_types={len(state.get('discoveredTaskTypes') or [])} eligible={len(eligible)} "
            f"errors={len(errors)}"
        )
        return state

    def eligible_gpus(self, requested: list[str] | None = None) -> list[str]:
        catalog = self.state.get("gpuCatalog") or {}
        names = list(requested or catalog.keys())
        return [
            gpu
            for gpu in names
            if gpu in catalog
            and (catalog[gpu].get("canVerify") is not False)
            and int(catalog[gpu].get("maxConcurrentTasks") or 0) > 0
        ]

    def task_types_for(self, client: Any, *, model_address: str, model_last_modified: Any, gpu: str) -> list[str]:
        now = utc_now()
        key = f"{model_address}|{model_last_modified or ''}|{gpu}"
        with self._lock:
            cache = self.state.setdefault("modelGpuTaskTypes", {})
            item = dict(cache.get(key) or {})
        ttl = 86400 if item.get("taskTypes") else 3600
        if _fresh(item.get("updatedAt"), now, ttl):
            return list(item.get("taskTypes") or [])
        try:
            payload = client.list_model_task_types(gpu, model_address)
            task_types = _extract_task_names(payload)
        except Exception as exc:
            if _fresh(item.get("updatedAt"), now, 86400):
                return list(item.get("taskTypes") or [])
            raise OfficialCapabilityUnavailable(
                f"model+GPU task capability unavailable for {gpu}: {type(exc).__name__}: {exc}"
            ) from exc
        with self._lock:
            cache = self.state.setdefault("modelGpuTaskTypes", {})
            cache[key] = {"updatedAt": now.isoformat(), "taskTypes": task_types}
            self._prune_model_gpu_cache(self.state)
            self.state["generatedAt"] = now.isoformat()
            write_json(self.path, self.state)
        return task_types

    def summary(self) -> dict[str, Any]:
        return {
            "enabled": True,
            "ready": self.ready,
            "pauseReason": self.pause_reason,
            "catalogUpdatedAt": self.state.get("catalogUpdatedAt"),
            "taskTreeUpdatedAt": self.state.get("taskTreeUpdatedAt"),
            "gpuCount": len(self.state.get("gpuCatalog") or {}),
            "discoveredTaskTypeCount": len(self.state.get("discoveredTaskTypes") or []),
            "errors": list(self.state.get("errors") or []),
        }