feat: add durable success-first modelhub agent
This commit is contained in:
CoolBoy
192
modelhub_submmit_api/routing_engine.py
Normal file
192
modelhub_submmit_api/routing_engine.py
Normal file
@@ -0,0 +1,192 @@
|
||||
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():
|
||||
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 {}),
|
||||
}
|
||||
Reference in New Issue
Block a user