submmit/modelhub_submmit_api/routing_engine.py

from __future__ import annotations

import math
from collections import defaultdict
from pathlib import Path
from typing import Any

from common import read_json, utc_now, write_json
from gpu_strategy import wilson_lower_bound
from submission_claims import candidate_key


ROUTING_STATE_VERSION = 1
DEFAULT_ROUTING_STATE_PATH = Path(".modelhub_state/routing_intelligence.json")


class SuccessFirstRoutingEngine:
    """Deterministic success-first routing; queue speed only breaks close races."""

    def __init__(
        self,
        path: Path | str = DEFAULT_ROUTING_STATE_PATH,
        *,
        outcome_stats: dict[str, Any] | None = None,
        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._load()
        self.outcome_stats = outcome_stats if isinstance(outcome_stats, dict) else {}

    def _load(self) -> dict[str, Any]:
        try:
            value = read_json(self.path)
        except (FileNotFoundError, ValueError):
            value = {}
        if not isinstance(value, dict) or int(value.get("version") or 0) != ROUTING_STATE_VERSION:
            value = {
                "version": ROUTING_STATE_VERSION,
                "generatedAt": utc_now().isoformat(),
                "acceptedTotal": 0,
                "acceptedSinceRefresh": 0,
                "acceptedByCategory": {"unified_success_first": 0},
                "acceptedByRoute": {},
            }
        return value

    def _evidence(self, candidate: dict[str, Any]) -> tuple[str, int, int, float]:
        profile = candidate.get("preflightMetadata") or {}
        gpu = str(candidate.get("targetGpu") or "")
        framework = str(candidate.get("framework") or "")
        task_type = str(candidate.get("taskType") or "")
        model_type = str(profile.get("modelType") or "").strip()
        quantization = str(profile.get("quantizationMethod") or "none").strip()
        try:
            load_bytes = max(1, int(profile.get("estimatedLoadBytes") or 0))
        except (TypeError, ValueError):
            load_bytes = 0
        if model_type and load_bytes > 0:
            size_bucket = int(load_bytes).bit_length() - 1
            key = f"{gpu}|{framework}|{task_type}|{model_type}|{quantization}|{size_bucket}"
            item = (self.outcome_stats.get("sizedProfileCombinationStats") or {}).get(key) or {}
            total = int(item.get("decisionTotal") or 0)
            if total >= 20:
                success = int(item.get("successCount") or 0)
                return "local_profile_size", success, total, wilson_lower_bound(success, total)
        if model_type:
            key = f"{gpu}|{framework}|{task_type}|{model_type}|{quantization}"
            item = (self.outcome_stats.get("profileCombinationStats") or {}).get(key) or {}
            total = int(item.get("decisionTotal") or 0)
            if total >= 40:
                success = int(item.get("successCount") or 0)
                return "local_profile", success, total, wilson_lower_bound(success, total)
        public_samples = max(0, int(candidate.get("frameworkMarketSamples") or 0))
        public_rate = candidate.get("frameworkMarketSuccessRate")
        if public_rate is None:
            public_success = 0
        else:
            public_success = max(0, min(public_samples, round(float(public_rate) * public_samples)))
        local_samples = max(0, int(candidate.get("frameworkLocalSamples") or 0))
        local_rate = candidate.get("frameworkLocalSuccessRate")
        local_success = 0 if local_rate is None else max(0, min(local_samples, round(float(local_rate) * local_samples)))

        if local_samples >= 100:
            total = local_samples
            success = local_success
            level = "local_task_gpu_framework"
        elif public_samples >= 300:
            total = public_samples
            success = public_success
            level = "official_task_gpu_framework"
        elif public_samples > 0:
            total = public_samples
            success = public_success
            level = "official_low_sample"
        else:
            total = 2
            success = 1
            level = "global_conservative_prior"
        return level, success, total, wilson_lower_bound(success, total)

    def annotate(self, candidate: dict[str, Any]) -> dict[str, Any]:
        annotated = dict(candidate)
        evidence, success, total, lower = self._evidence(candidate)
        raw_eta = candidate.get("queueBacklogHours")
        try:
            eta = max(0.25, float(raw_eta))
            if not math.isfinite(eta) or eta >= 9999:
                raise ValueError
            eta_unknown = False
        except (TypeError, ValueError):
            eta = 36.0
            eta_unknown = True
        success_band = int(lower / 0.05)
        expected = (lower * lower) / eta
        annotated.update(
            {
                "routingEvidenceLevel": evidence,
                "routingSuccesses": success,
                "routingSamples": total,
                "routingSuccessLowerBound": lower,
                "routingSuccessBand": success_band,
                "routingEtaHours": eta,
                "routingEtaUnknown": eta_unknown,
                "routingExpectedSuccessPerHour": expected,
                "strategyCategory": "unified_success_first",
            }
        )
        return annotated

    def order_candidates(self, candidates: list[dict[str, Any]]) -> list[dict[str, Any]]:
        if not candidates:
            return []
        annotated = [self.annotate(candidate) for candidate in candidates]
        by_model: dict[str, list[dict[str, Any]]] = defaultdict(list)
        for candidate in annotated:
            by_model[str(candidate.get("repoId") or candidate.get("modelAddress") or "")].append(candidate)

        selected: list[dict[str, Any]] = []
        for routes in by_model.values():
            # The platform permits some transformers submissions only after a
            # non-transformers attempt has failed.  Prefer an immediately
            # executable route whenever one exists; transformers remains a
            # fallback for tasks/models with no other official route.
            non_transformers = [
                item for item in routes if str(item.get("framework") or "").casefold() != "transformers"
            ]
            transformers_unlocked = any(
                bool(item.get("transformersPrerequisiteSatisfied"))
                for item in routes
                if str(item.get("framework") or "").casefold() == "transformers"
            )
            if non_transformers and not transformers_unlocked:
                routes = non_transformers
            best = max(float(item["routingSuccessLowerBound"]) for item in routes)
            close = [item for item in routes if best - float(item["routingSuccessLowerBound"]) <= 0.05]
            close.sort(
                key=lambda item: (
                    -int(item["routingSuccessBand"]),
                    -float(item["routingExpectedSuccessPerHour"]),
                    -int(item["routingSamples"]),
                    str(item.get("targetGpu") or ""),
                    str(item.get("framework") or ""),
                )
            )
            selected.append(close[0])

        selected.sort(
            key=lambda item: (
                -int(item["routingSuccessBand"]),
                -float(item["routingExpectedSuccessPerHour"]),
                -int(item["routingSamples"]),
                -int(item.get("downloads") or 0),
                str(item.get("repoId") or ""),
            )
        )
        return selected

    @property
    def submissions_until_refresh(self) -> int:
        return 1_000_000_000

    def record_accepted(self, candidates: list[dict[str, Any]]) -> dict[str, Any]:
        routes = self.state.setdefault("acceptedByRoute", {})
        for candidate in candidates:
            key = "|".join(
                str(candidate.get(name) or "")
                for name in ("taskType", "targetGpu", "framework")
            )
            routes[key] = int(routes.get(key) or 0) + 1
        self.state["acceptedTotal"] = int(self.state.get("acceptedTotal") or 0) + len(candidates)
        self.state["acceptedSinceRefresh"] = int(self.state.get("acceptedSinceRefresh") or 0) + len(candidates)
        categories = self.state.setdefault("acceptedByCategory", {"unified_success_first": 0})
        categories["unified_success_first"] = int(categories.get("unified_success_first") or 0) + len(candidates)
        self.state["generatedAt"] = utc_now().isoformat()
        write_json(self.path, self.state)
        return self.state

    def summary(self) -> dict[str, Any]:
        return {
            "enabled": True,
            "mode": "unified_success_first",
            "statePath": str(self.path),
            "acceptedTotal": int(self.state.get("acceptedTotal") or 0),
            "routeCount": len(self.state.get("acceptedByRoute") or {}),
        }