feat: add adaptive GPU scheduling
This commit is contained in:
CoolBoy
480
modelhub_submmit_api/gpu_strategy.py
Normal file
480
modelhub_submmit_api/gpu_strategy.py
Normal file
@@ -0,0 +1,480 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections import defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
from common import parse_datetime, read_json, utc_now, write_json
|
||||
from history_stats import is_failure, is_success
|
||||
from modelhub_client import ModelHubClientPool
|
||||
from submission_claims import candidate_key
|
||||
|
||||
|
||||
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
|
||||
|
||||
LONG_TERM = "long_term"
|
||||
ALL_SUPPORTED = "all_supported"
|
||||
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 _gpu_name(record: dict[str, Any]) -> str:
|
||||
return str(record.get("gpuType") or record.get("targetGpu") or "").strip()
|
||||
|
||||
|
||||
def _terminal_outcome(record: dict[str, Any]) -> str | None:
|
||||
try:
|
||||
if is_success(record):
|
||||
return "success"
|
||||
if is_failure(record):
|
||||
return "failure"
|
||||
except TypeError:
|
||||
# Be tolerant of APIs that serialize verifyResult as a string.
|
||||
copied = dict(record)
|
||||
try:
|
||||
copied["verifyResult"] = float(record.get("verifyResult"))
|
||||
except (TypeError, ValueError):
|
||||
copied["verifyResult"] = None
|
||||
if is_success(copied):
|
||||
return "success"
|
||||
if is_failure(copied):
|
||||
return "failure"
|
||||
return None
|
||||
|
||||
|
||||
def _task_sort_key(record: dict[str, Any]) -> tuple[float, str]:
|
||||
timestamp = (
|
||||
parse_datetime(record.get("updateTime"))
|
||||
or parse_datetime(record.get("createTime"))
|
||||
or parse_datetime(record.get("submitTime"))
|
||||
)
|
||||
return (timestamp.timestamp() if timestamp else 0.0, str(record.get("taskId") or ""))
|
||||
|
||||
|
||||
def wilson_lower_bound(successes: int, total: int, *, z: float = 1.96) -> float:
|
||||
if total <= 0:
|
||||
return 0.0
|
||||
probability = successes / total
|
||||
z_squared = z * z
|
||||
denominator = 1.0 + z_squared / total
|
||||
centre = probability + z_squared / (2.0 * total)
|
||||
margin = z * math.sqrt(
|
||||
(probability * (1.0 - probability) / total) + z_squared / (4.0 * total * total)
|
||||
)
|
||||
return max(0.0, (centre - margin) / denominator)
|
||||
|
||||
|
||||
def _summarize_gpu_records(
|
||||
records: list[dict[str, Any]],
|
||||
supported_gpus: list[str],
|
||||
) -> dict[str, dict[str, Any]]:
|
||||
supported = set(supported_gpus)
|
||||
counts: dict[str, dict[str, int]] = {
|
||||
gpu: {"success": 0, "failure": 0}
|
||||
for gpu in supported_gpus
|
||||
}
|
||||
for record in records:
|
||||
gpu = _gpu_name(record)
|
||||
if gpu not in supported:
|
||||
continue
|
||||
outcome = _terminal_outcome(record)
|
||||
if outcome is not None:
|
||||
counts[gpu][outcome] += 1
|
||||
|
||||
summaries: dict[str, dict[str, Any]] = {}
|
||||
for gpu in supported_gpus:
|
||||
success = counts[gpu]["success"]
|
||||
failure = counts[gpu]["failure"]
|
||||
terminal = success + failure
|
||||
summaries[gpu] = {
|
||||
"gpu": gpu,
|
||||
"success": success,
|
||||
"failure": failure,
|
||||
"terminal": terminal,
|
||||
"successRate": success / terminal if terminal else 0.0,
|
||||
"wilsonLowerBound": wilson_lower_bound(success, terminal),
|
||||
}
|
||||
return summaries
|
||||
|
||||
|
||||
def _rank_gpu_summaries(summaries: dict[str, dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
return sorted(
|
||||
summaries.values(),
|
||||
key=lambda item: (
|
||||
-float(item["wilsonLowerBound"]),
|
||||
-float(item["successRate"]),
|
||||
-int(item["terminal"]),
|
||||
str(item["gpu"]),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_strategy_snapshot(
|
||||
tasks: list[dict[str, Any]],
|
||||
*,
|
||||
supported_gpus: list[str],
|
||||
generated_at: datetime | None = None,
|
||||
generation: int = 0,
|
||||
recent_terminal_window: int = DEFAULT_RECENT_TERMINAL_WINDOW,
|
||||
long_term_min_samples: int = DEFAULT_LONG_TERM_MIN_SAMPLES,
|
||||
refresh_submissions: int = DEFAULT_REFRESH_SUBMISSIONS,
|
||||
) -> dict[str, Any]:
|
||||
generated_at = generated_at or utc_now()
|
||||
supported_gpus = list(dict.fromkeys(gpu for gpu in supported_gpus if gpu))
|
||||
if not supported_gpus:
|
||||
raise ValueError("At least one supported GPU is required")
|
||||
|
||||
all_time = _summarize_gpu_records(tasks, supported_gpus)
|
||||
all_time_ranked = _rank_gpu_summaries(all_time)
|
||||
qualified_long = [
|
||||
item for item in all_time_ranked
|
||||
if int(item["terminal"]) >= max(1, int(long_term_min_samples))
|
||||
]
|
||||
long_term_gpus = [str(item["gpu"]) for item in qualified_long[:3]]
|
||||
if len(long_term_gpus) < min(3, len(supported_gpus)):
|
||||
for item in all_time_ranked:
|
||||
gpu = str(item["gpu"])
|
||||
if gpu not in long_term_gpus:
|
||||
long_term_gpus.append(gpu)
|
||||
if len(long_term_gpus) >= min(3, len(supported_gpus)):
|
||||
break
|
||||
|
||||
terminal_tasks = [record for record in tasks if _terminal_outcome(record) is not None]
|
||||
terminal_tasks.sort(key=_task_sort_key, reverse=True)
|
||||
recent_tasks = terminal_tasks[: max(1, int(recent_terminal_window))]
|
||||
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]
|
||||
|
||||
return {
|
||||
"version": STRATEGY_STATE_VERSION,
|
||||
"generation": max(0, int(generation)),
|
||||
"generatedAt": generated_at.isoformat(),
|
||||
"historyReady": True,
|
||||
"historyTaskCount": len(tasks),
|
||||
"supportedGpus": supported_gpus,
|
||||
"longTermGpus": long_term_gpus,
|
||||
"recentGpu": recent_gpu,
|
||||
"recentTerminalCount": len(recent_tasks),
|
||||
"refreshSubmissions": max(1, int(refresh_submissions)),
|
||||
"recentTerminalWindow": max(1, int(recent_terminal_window)),
|
||||
"longTermMinSamples": max(1, int(long_term_min_samples)),
|
||||
"acceptedSinceRefresh": 0,
|
||||
"acceptedTotal": 0,
|
||||
"acceptedByCategory": _empty_category_counts(),
|
||||
"longTermStats": all_time_ranked,
|
||||
"recentStats": recent_ranked,
|
||||
"lastRefreshError": None,
|
||||
"refreshRetryAfter": None,
|
||||
}
|
||||
|
||||
|
||||
def choose_next_category(counts: dict[str, int]) -> str:
|
||||
normalized = {category: max(0, int(counts.get(category, 0))) for category in CATEGORIES}
|
||||
next_total = sum(normalized.values()) + 1
|
||||
|
||||
def deficit(category: str) -> int:
|
||||
return CATEGORY_WEIGHTS[category] * next_total - normalized[category] * sum(CATEGORY_WEIGHTS.values())
|
||||
|
||||
return max(CATEGORIES, key=lambda category: (deficit(category), -CATEGORIES.index(category)))
|
||||
|
||||
|
||||
class GPUStrategyManager:
|
||||
def __init__(
|
||||
self,
|
||||
path: Path | str = DEFAULT_GPU_STRATEGY_PATH,
|
||||
*,
|
||||
refresh_submissions: int = DEFAULT_REFRESH_SUBMISSIONS,
|
||||
recent_terminal_window: int = DEFAULT_RECENT_TERMINAL_WINDOW,
|
||||
long_term_min_samples: int = DEFAULT_LONG_TERM_MIN_SAMPLES,
|
||||
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.log = log_fn or (lambda message: print(message, flush=True))
|
||||
self.state: dict[str, Any] | None = None
|
||||
|
||||
def _load(self) -> dict[str, Any] | None:
|
||||
try:
|
||||
value = read_json(self.path)
|
||||
except (FileNotFoundError, ValueError):
|
||||
return None
|
||||
return value if isinstance(value, dict) else None
|
||||
|
||||
def _is_compatible(self, state: dict[str, Any] | None, supported_gpus: list[str]) -> bool:
|
||||
if not state or int(state.get("version") or 0) != STRATEGY_STATE_VERSION:
|
||||
return False
|
||||
return list(state.get("supportedGpus") or []) == supported_gpus
|
||||
|
||||
def _refresh_due(self, state: dict[str, Any] | None, supported_gpus: list[str], now: datetime) -> tuple[bool, str]:
|
||||
if not self._is_compatible(state, supported_gpus):
|
||||
return True, "initial_or_gpu_catalog_changed"
|
||||
retry_after = parse_datetime(state.get("refreshRetryAfter")) if state is not None else None
|
||||
if state is not None and state.get("lastRefreshError") and retry_after is not None and retry_after > now:
|
||||
return False, "refresh_error_backoff"
|
||||
if not bool(state.get("historyReady", False)):
|
||||
return (retry_after is None or retry_after <= now), "history_not_ready"
|
||||
accepted = int(state.get("acceptedSinceRefresh") or 0)
|
||||
refresh_every = max(1, int(state.get("refreshSubmissions") or self.refresh_submissions))
|
||||
return accepted >= refresh_every, "accepted_submission_threshold"
|
||||
|
||||
@staticmethod
|
||||
def _load_platform_history(client: Any) -> list[dict[str, Any]]:
|
||||
if isinstance(client, ModelHubClientPool):
|
||||
by_account: dict[int, list[dict[str, Any]]] = {}
|
||||
errors: list[int] = []
|
||||
with ThreadPoolExecutor(max_workers=min(len(client.clients), 12)) as executor:
|
||||
futures = {
|
||||
executor.submit(account_client.list_tasks, page_size=100, only_mine=True): index
|
||||
for index, account_client in enumerate(client.clients, start=1)
|
||||
}
|
||||
for future in as_completed(futures):
|
||||
index = futures[future]
|
||||
try:
|
||||
by_account[index] = future.result()
|
||||
except Exception:
|
||||
errors.append(index)
|
||||
if errors:
|
||||
joined = ",".join(str(index) for index in sorted(errors))
|
||||
raise RuntimeError(f"ModelHub history fetch failed for account indexes: {joined}")
|
||||
tasks = [task for index in sorted(by_account) for task in by_account[index]]
|
||||
else:
|
||||
tasks = client.list_tasks(page_size=100, only_mine=True)
|
||||
|
||||
deduped: list[dict[str, Any]] = []
|
||||
seen_task_ids: set[str] = set()
|
||||
for task in tasks:
|
||||
if not isinstance(task, dict):
|
||||
continue
|
||||
task_id = str(task.get("taskId")) if task.get("taskId") is not None else None
|
||||
if task_id and task_id in seen_task_ids:
|
||||
continue
|
||||
if task_id:
|
||||
seen_task_ids.add(task_id)
|
||||
deduped.append(task)
|
||||
return deduped
|
||||
|
||||
def prepare(
|
||||
self,
|
||||
client: Any,
|
||||
*,
|
||||
supported_gpus: list[str],
|
||||
now: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
now = now or utc_now()
|
||||
supported_gpus = list(dict.fromkeys(gpu for gpu in supported_gpus if gpu))
|
||||
state = self._load()
|
||||
refresh_due, reason = self._refresh_due(state, supported_gpus, now)
|
||||
if not refresh_due and state is not None:
|
||||
self.state = state
|
||||
self._log_state("loaded")
|
||||
return state
|
||||
|
||||
previous_generation = int(state.get("generation", -1)) if state is not None else -1
|
||||
self.log(f"[strategy] refresh_start reason={reason} supported_gpus={len(supported_gpus)}")
|
||||
try:
|
||||
tasks = self._load_platform_history(client)
|
||||
refreshed = build_strategy_snapshot(
|
||||
tasks,
|
||||
supported_gpus=supported_gpus,
|
||||
generated_at=now,
|
||||
generation=previous_generation + 1,
|
||||
recent_terminal_window=self.recent_terminal_window,
|
||||
long_term_min_samples=self.long_term_min_samples,
|
||||
refresh_submissions=self.refresh_submissions,
|
||||
)
|
||||
refreshed["acceptedTotal"] = int((state or {}).get("acceptedTotal") or 0)
|
||||
self.state = refreshed
|
||||
write_json(self.path, refreshed)
|
||||
self._log_state("refreshed")
|
||||
return refreshed
|
||||
except Exception as exc:
|
||||
self.log(f"[strategy] refresh_error reason={type(exc).__name__}: {exc}")
|
||||
if self._is_compatible(state, supported_gpus):
|
||||
assert state is not None
|
||||
state["lastRefreshError"] = str(exc)
|
||||
state["refreshRetryAfter"] = (now + timedelta(minutes=5)).isoformat()
|
||||
self.state = state
|
||||
write_json(self.path, state)
|
||||
return state
|
||||
|
||||
fallback = build_strategy_snapshot(
|
||||
[],
|
||||
supported_gpus=supported_gpus,
|
||||
generated_at=now,
|
||||
generation=0,
|
||||
recent_terminal_window=self.recent_terminal_window,
|
||||
long_term_min_samples=self.long_term_min_samples,
|
||||
refresh_submissions=self.refresh_submissions,
|
||||
)
|
||||
fallback["historyReady"] = False
|
||||
fallback["lastRefreshError"] = str(exc)
|
||||
fallback["refreshRetryAfter"] = (now + timedelta(minutes=5)).isoformat()
|
||||
self.state = fallback
|
||||
write_json(self.path, fallback)
|
||||
self._log_state("fallback")
|
||||
return fallback
|
||||
|
||||
def _log_state(self, action: str) -> None:
|
||||
if self.state is None:
|
||||
return
|
||||
counts = self.state.get("acceptedByCategory") or {}
|
||||
self.log(
|
||||
f"[strategy] {action} generation={self.state.get('generation', 0)} "
|
||||
f"accepted={self.state.get('acceptedSinceRefresh', 0)}/{self.state.get('refreshSubmissions', self.refresh_submissions)} "
|
||||
f"categories={counts.get(LONG_TERM, 0)},{counts.get(ALL_SUPPORTED, 0)},{counts.get(RECENT, 0)} "
|
||||
f"long={','.join(self.state.get('longTermGpus') or [])} recent={self.state.get('recentGpu') or 'n/a'}"
|
||||
)
|
||||
|
||||
@property
|
||||
def submissions_until_refresh(self) -> int:
|
||||
if self.state is None:
|
||||
return self.refresh_submissions
|
||||
refresh_every = max(1, int(self.state.get("refreshSubmissions") or self.refresh_submissions))
|
||||
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]:
|
||||
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 "")
|
||||
|
||||
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 []
|
||||
|
||||
def order_candidates(self, candidates: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
if self.state is None or not candidates:
|
||||
return list(candidates)
|
||||
|
||||
by_gpu: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for candidate in candidates:
|
||||
by_gpu[str(candidate.get("targetGpu") or "")].append(candidate)
|
||||
gpu_indexes: dict[str, int] = defaultdict(int)
|
||||
used: set[str] = set()
|
||||
ordered: list[dict[str, Any]] = []
|
||||
virtual_counts = {
|
||||
category: max(0, int((self.state.get("acceptedByCategory") or {}).get(category, 0)))
|
||||
for category in CATEGORIES
|
||||
}
|
||||
|
||||
def take_from_gpu(gpu: str) -> dict[str, Any] | None:
|
||||
pool = by_gpu.get(gpu) or []
|
||||
index = gpu_indexes[gpu]
|
||||
while index < len(pool):
|
||||
candidate = pool[index]
|
||||
index += 1
|
||||
gpu_indexes[gpu] = index
|
||||
if candidate_key(candidate) not in used:
|
||||
return candidate
|
||||
gpu_indexes[gpu] = index
|
||||
return None
|
||||
|
||||
while len(ordered) < len(candidates):
|
||||
planned_category = choose_next_category(virtual_counts)
|
||||
selected: dict[str, Any] | None = None
|
||||
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):
|
||||
selected = take_from_gpu(gpu)
|
||||
if selected is not None:
|
||||
actual_category = category
|
||||
break
|
||||
if selected is not None:
|
||||
break
|
||||
|
||||
if selected is None:
|
||||
# Unknown/custom GPUs can only appear when callers bypass the normal resolver.
|
||||
selected = next((candidate for candidate in candidates if candidate_key(candidate) not in used), None)
|
||||
actual_category = ALL_SUPPORTED
|
||||
if selected is None:
|
||||
break
|
||||
|
||||
key = candidate_key(selected)
|
||||
used.add(key)
|
||||
annotated = dict(selected)
|
||||
annotated["strategyCategory"] = actual_category
|
||||
annotated["strategyPlannedCategory"] = planned_category
|
||||
annotated["strategyGeneration"] = int(self.state.get("generation") or 0)
|
||||
ordered.append(annotated)
|
||||
virtual_counts[actual_category] += 1
|
||||
|
||||
return ordered
|
||||
|
||||
def record_accepted(self, candidates: list[dict[str, Any]]) -> dict[str, Any] | None:
|
||||
if not candidates:
|
||||
return self.state
|
||||
state = self._load() or self.state
|
||||
if state is None:
|
||||
return None
|
||||
|
||||
category_counts = {
|
||||
category: max(0, int((state.get("acceptedByCategory") or {}).get(category, 0)))
|
||||
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
|
||||
|
||||
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
|
||||
self.state = state
|
||||
write_json(self.path, state)
|
||||
self._log_state("progress")
|
||||
return state
|
||||
|
||||
def summary(self) -> dict[str, Any]:
|
||||
if self.state is None:
|
||||
return {"enabled": False}
|
||||
return {
|
||||
"enabled": True,
|
||||
"statePath": str(self.path),
|
||||
"generation": int(self.state.get("generation") or 0),
|
||||
"generatedAt": self.state.get("generatedAt"),
|
||||
"historyReady": bool(self.state.get("historyReady", False)),
|
||||
"historyTaskCount": int(self.state.get("historyTaskCount") or 0),
|
||||
"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 {}),
|
||||
"longTermGpus": list(self.state.get("longTermGpus") or []),
|
||||
"recentGpu": self.state.get("recentGpu"),
|
||||
"recentTerminalCount": int(self.state.get("recentTerminalCount") or 0),
|
||||
"refreshDue": self.submissions_until_refresh <= 0,
|
||||
}
|
||||
Reference in New Issue
Block a user