feat: add queue-aware adaptive scheduling

This commit is contained in:
CoolBoy
2026-08-04 20:22:08 +08:00
parent ccae7ff8f3
commit 882479e43e
12 changed files with 1586 additions and 95 deletions

View File

@@ -17,7 +17,7 @@ DEFAULT_GPU_STRATEGY_PATH = Path(".modelhub_state/gpu_strategy.json")
DEFAULT_REFRESH_SUBMISSIONS = 200
DEFAULT_RECENT_TERMINAL_WINDOW = 1000
DEFAULT_LONG_TERM_MIN_SAMPLES = 100
STRATEGY_STATE_VERSION = 1
STRATEGY_STATE_VERSION = 2
LONG_TERM = "long_term"
ALL_SUPPORTED = "all_supported"
@@ -25,14 +25,17 @@ RECENT = "recent"
CATEGORIES = (LONG_TERM, ALL_SUPPORTED, RECENT)
CATEGORY_WEIGHTS = {LONG_TERM: 5, ALL_SUPPORTED: 3, RECENT: 2}
# The long-term half is split 50/30/20 across its three ranked GPUs.
LONG_TERM_GPU_PATTERN = (0, 1, 0, 2, 0, 1, 0, 1, 0, 2)
def _empty_category_counts() -> dict[str, int]:
return {category: 0 for category in CATEGORIES}
def _empty_gpu_category_counts(supported_gpus: list[str]) -> dict[str, dict[str, int]]:
return {
category: {gpu: 0 for gpu in supported_gpus}
for category in CATEGORIES
}
def _gpu_name(record: dict[str, Any]) -> str:
return str(record.get("gpuType") or record.get("targetGpu") or "").strip()
@@ -160,12 +163,11 @@ def build_strategy_snapshot(
recent = _summarize_gpu_records(recent_tasks, supported_gpus)
recent_ranked = _rank_gpu_summaries(recent)
recent_qualified = [item for item in recent_ranked if int(item["terminal"]) >= 20]
if recent_qualified:
recent_gpu = str(recent_qualified[0]["gpu"])
elif recent_tasks and recent_ranked:
recent_gpu = str(recent_ranked[0]["gpu"])
else:
recent_gpu = long_term_gpus[0]
recent_source = recent_qualified if recent_qualified else (recent_ranked if recent_tasks else [])
recent_gpus = [str(item["gpu"]) for item in recent_source[: min(3, len(supported_gpus))]]
if not recent_gpus:
recent_gpus = list(long_term_gpus[: min(3, len(supported_gpus))])
recent_gpu = recent_gpus[0]
return {
"version": STRATEGY_STATE_VERSION,
@@ -176,6 +178,7 @@ def build_strategy_snapshot(
"supportedGpus": supported_gpus,
"longTermGpus": long_term_gpus,
"recentGpu": recent_gpu,
"recentGpus": recent_gpus,
"recentTerminalCount": len(recent_tasks),
"refreshSubmissions": max(1, int(refresh_submissions)),
"recentTerminalWindow": max(1, int(recent_terminal_window)),
@@ -183,6 +186,7 @@ def build_strategy_snapshot(
"acceptedSinceRefresh": 0,
"acceptedTotal": 0,
"acceptedByCategory": _empty_category_counts(),
"acceptedByGpuCategory": _empty_gpu_category_counts(supported_gpus),
"longTermStats": all_time_ranked,
"recentStats": recent_ranked,
"lastRefreshError": None,
@@ -208,15 +212,20 @@ class GPUStrategyManager:
refresh_submissions: int = DEFAULT_REFRESH_SUBMISSIONS,
recent_terminal_window: int = DEFAULT_RECENT_TERMINAL_WINDOW,
long_term_min_samples: int = DEFAULT_LONG_TERM_MIN_SAMPLES,
market_intelligence: Any | None = None,
log_fn: Callable[[str], None] | None = None,
) -> None:
self.path = Path(path)
self.refresh_submissions = max(1, int(refresh_submissions))
self.recent_terminal_window = max(1, int(recent_terminal_window))
self.long_term_min_samples = max(1, int(long_term_min_samples))
self.market_intelligence = market_intelligence
self.log = log_fn or (lambda message: print(message, flush=True))
self.state: dict[str, Any] | None = None
def set_market_intelligence(self, market_intelligence: Any | None) -> None:
self.market_intelligence = market_intelligence
def _load(self) -> dict[str, Any] | None:
try:
value = read_json(self.path)
@@ -357,22 +366,49 @@ class GPUStrategyManager:
accepted = max(0, int(self.state.get("acceptedSinceRefresh") or 0))
return max(0, refresh_every - accepted)
def _category_gpu_order(self, category: str, occurrence: int) -> list[str]:
def _category_gpu_order(self, category: str, gpu_counts: dict[str, int]) -> list[str]:
assert self.state is not None
supported = list(self.state.get("supportedGpus") or [])
long_term = [gpu for gpu in self.state.get("longTermGpus") or [] if gpu in supported]
recent_gpu = str(self.state.get("recentGpu") or "")
recent = [gpu for gpu in self.state.get("recentGpus") or [] if gpu in supported]
if not recent:
recent_gpu = str(self.state.get("recentGpu") or "")
if recent_gpu in supported:
recent = [recent_gpu]
if category == LONG_TERM and long_term:
desired_index = LONG_TERM_GPU_PATTERN[occurrence % len(LONG_TERM_GPU_PATTERN)]
desired = long_term[min(desired_index, len(long_term) - 1)]
return [desired, *(gpu for gpu in long_term if gpu != desired)]
if category == RECENT and recent_gpu in supported:
return [recent_gpu]
if supported:
offset = occurrence % len(supported)
return [*supported[offset:], *supported[:offset]]
return []
if category == LONG_TERM:
eligible = long_term
base_pattern = (5.0, 3.0, 2.0)
elif category == RECENT:
eligible = recent
base_pattern = (6.0, 3.0, 1.0)
else:
eligible = supported
base_pattern = tuple(1.0 for _ in eligible)
if not eligible:
return []
weights: dict[str, float] = {}
for index, gpu in enumerate(eligible):
base = base_pattern[min(index, len(base_pattern) - 1)]
market_weight = 1.0
if self.market_intelligence is not None:
try:
market_weight = float(self.market_intelligence.gpu_weight(gpu, category=category))
except Exception:
market_weight = 1.0
weights[gpu] = max(0.05, base * market_weight)
total_weight = sum(weights.values()) or 1.0
next_total = sum(max(0, int(gpu_counts.get(gpu, 0))) for gpu in eligible) + 1
rank_index = {gpu: index for index, gpu in enumerate(eligible)}
return sorted(
eligible,
key=lambda gpu: (
-((weights[gpu] / total_weight) * next_total - max(0, int(gpu_counts.get(gpu, 0)))),
rank_index[gpu],
),
)
def order_candidates(self, candidates: list[dict[str, Any]]) -> list[dict[str, Any]]:
if self.state is None or not candidates:
@@ -388,6 +424,14 @@ class GPUStrategyManager:
category: max(0, int((self.state.get("acceptedByCategory") or {}).get(category, 0)))
for category in CATEGORIES
}
stored_gpu_counts = self.state.get("acceptedByGpuCategory") or {}
virtual_gpu_counts = {
category: {
gpu: max(0, int((stored_gpu_counts.get(category) or {}).get(gpu, 0)))
for gpu in self.state.get("supportedGpus") or []
}
for category in CATEGORIES
}
def take_from_gpu(gpu: str) -> dict[str, Any] | None:
pool = by_gpu.get(gpu) or []
@@ -407,8 +451,7 @@ class GPUStrategyManager:
actual_category = planned_category
categories_to_try = [planned_category, *(category for category in CATEGORIES if category != planned_category)]
for category in categories_to_try:
occurrence = virtual_counts[category]
for gpu in self._category_gpu_order(category, occurrence):
for gpu in self._category_gpu_order(category, virtual_gpu_counts[category]):
selected = take_from_gpu(gpu)
if selected is not None:
actual_category = category
@@ -429,8 +472,18 @@ class GPUStrategyManager:
annotated["strategyCategory"] = actual_category
annotated["strategyPlannedCategory"] = planned_category
annotated["strategyGeneration"] = int(self.state.get("generation") or 0)
if self.market_intelligence is not None:
try:
annotated.update(self.market_intelligence.gpu_metadata(str(annotated.get("targetGpu") or "")))
except Exception:
pass
ordered.append(annotated)
virtual_counts[actual_category] += 1
selected_gpu = str(annotated.get("targetGpu") or "")
if selected_gpu:
virtual_gpu_counts[actual_category][selected_gpu] = (
virtual_gpu_counts[actual_category].get(selected_gpu, 0) + 1
)
return ordered
@@ -445,16 +498,29 @@ class GPUStrategyManager:
category: max(0, int((state.get("acceptedByCategory") or {}).get(category, 0)))
for category in CATEGORIES
}
supported = list(state.get("supportedGpus") or [])
stored_gpu_counts = state.get("acceptedByGpuCategory") or {}
gpu_category_counts = {
category: {
gpu: max(0, int((stored_gpu_counts.get(category) or {}).get(gpu, 0)))
for gpu in supported
}
for category in CATEGORIES
}
for candidate in candidates:
category = str(candidate.get("strategyCategory") or ALL_SUPPORTED)
if category not in category_counts:
category = ALL_SUPPORTED
category_counts[category] += 1
gpu = str(candidate.get("targetGpu") or "")
if gpu in gpu_category_counts[category]:
gpu_category_counts[category][gpu] += 1
accepted_count = len(candidates)
state["acceptedSinceRefresh"] = int(state.get("acceptedSinceRefresh") or 0) + accepted_count
state["acceptedTotal"] = int(state.get("acceptedTotal") or 0) + accepted_count
state["acceptedByCategory"] = category_counts
state["acceptedByGpuCategory"] = gpu_category_counts
self.state = state
write_json(self.path, state)
self._log_state("progress")
@@ -473,8 +539,10 @@ class GPUStrategyManager:
"acceptedSinceRefresh": int(self.state.get("acceptedSinceRefresh") or 0),
"refreshSubmissions": int(self.state.get("refreshSubmissions") or self.refresh_submissions),
"acceptedByCategory": dict(self.state.get("acceptedByCategory") or {}),
"acceptedByGpuCategory": dict(self.state.get("acceptedByGpuCategory") or {}),
"longTermGpus": list(self.state.get("longTermGpus") or []),
"recentGpu": self.state.get("recentGpu"),
"recentGpus": list(self.state.get("recentGpus") or []),
"recentTerminalCount": int(self.state.get("recentTerminalCount") or 0),
"refreshDue": self.submissions_until_refresh <= 0,
}