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 {}), }