feat: add queue-aware adaptive scheduling
This commit is contained in:
@@ -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,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user