207 lines
8.9 KiB
Python
207 lines
8.9 KiB
Python
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 {}),
|
|
}
|