CoolBoy/submmit
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat: reserve queue capacity for recent models

Browse Source
This commit is contained in:
CoolBoy
2026-08-11 01:29:54 +08:00
parent 38ab25fc3c
commit f12d96b138
13 changed files with 927 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

261
modelhub_submmit_api/queue_cleanup.py
View File

@@ -3,11 +3,12 @@ from __future__ import annotations
import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Callable, Iterable
from candidate_preflight import CandidatePreflightAdvisor, MODEL_LOAD_OVERHEAD
from common import write_json
from common import utc_now, write_json
from hf_discovery import HuggingFaceDiscovery, inspect_repo_tree
from modelhub_client import ModelHubClient, ModelHubClientPool
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
@@ -159,6 +160,41 @@ def _load_repository_sizes(
return sizes, errors
def _load_model_last_modified(
model_ids: set[str],
*,
discovery: HuggingFaceDiscovery,
read_concurrency: int,
log: Callable[[str], None],
) -> tuple[dict[str, datetime], dict[str, str]]:
values: dict[str, datetime] = {}
errors: dict[str, str] = {}
completed = 0
workers = min(max(1, int(read_concurrency)), max(1, len(model_ids)))
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = {
executor.submit(discovery.get_model_last_modified, model_id): model_id
for model_id in sorted(model_ids)
}
for future in as_completed(futures):
model_id = futures[future]
try:
value = future.result()
if value is None:
errors[model_id] = "model_last_modified_unknown"
else:
values[model_id] = value
except Exception as exc:
errors[model_id] = f"{type(exc).__name__}: {exc}"
completed += 1
if completed == len(model_ids) or completed % 50 == 0:
log(
f"[queue-cleanup] age_scan={completed}/{len(model_ids)} "
f"complete={len(values)} unknown={len(errors)}"
)
return values, errors
def find_certain_oom_tasks(
tasks: list[OwnedTask],
*,
@@ -202,6 +238,62 @@ def find_certain_oom_tasks(
return decisions, skipped
def find_old_overflow_tasks(
tasks: list[OwnedTask],
*,
model_last_modified: dict[str, datetime],
queue_threshold: int = 80,
recent_model_days: int = 7,
reference_time: datetime | None = None,
incomplete_accounts: set[int] | None = None,
) -> tuple[list[dict[str, Any]], dict[str, int]]:
threshold = max(1, int(queue_threshold))
recent_days = max(1, int(recent_model_days))
cutoff = (reference_time or utc_now()) - timedelta(days=recent_days)
incomplete_accounts = incomplete_accounts or set()
grouped: dict[int, list[OwnedTask]] = {}
for task in tasks:
grouped.setdefault(task.account_index, []).append(task)
decisions: list[dict[str, Any]] = []
skipped = {
"accountsWithIncompleteListing": len(incomplete_accounts),
"withinFirstQueuePositions": 0,
"recentOverflowTasks": 0,
"modelAgeUnknown": 0,
}
for account_index, account_tasks in sorted(grouped.items()):
if account_index in incomplete_accounts:
continue
ordered = sorted(account_tasks, key=lambda item: item.task_id)
skipped["withinFirstQueuePositions"] += min(threshold, len(ordered))
for position, task in enumerate(ordered, start=1):
if position <= threshold:
continue
last_modified = model_last_modified.get(task.model_id)
if last_modified is None:
skipped["modelAgeUnknown"] += 1
continue
if last_modified >= cutoff:
skipped["recentOverflowTasks"] += 1
continue
decisions.append(
{
"accountIndex": account_index + 1,
"taskId": task.task_id,
"modelId": task.model_id,
"gpuType": task.gpu_type,
"status": task.status,
"queuePosition": position,
"queueThreshold": threshold,
"modelLastModified": last_modified.isoformat(),
"recentCutoff": cutoff.isoformat(),
"reason": "old_model_beyond_account_queue_threshold",
}
)
return decisions, skipped
def _chunks(values: list[int], size: int) -> Iterable[list[int]]:
for offset in range(0, len(values), size):
yield values[offset : offset + size]
@@ -215,10 +307,20 @@ def cleanup_certain_oom_tasks(
read_concurrency: int = 6,
stop_batch_size: int = DEFAULT_STOP_BATCH_SIZE,
gpu_memory_gib: dict[str, float] | None = None,
age_queue_threshold: int | None = None,
reference_time: datetime | None = None,
log: Callable[[str], None] = print,
) -> dict[str, Any]:
"""Stop only active tasks that cannot fit the selected GPU by known capacity."""
"""Stop deterministic OOM tasks and old tasks beyond each account's protected prefix."""
clients = list(modelhub.clients)
reference_time = reference_time or utc_now()
configured_queue_threshold = (
age_queue_threshold
if age_queue_threshold is not None
else getattr(modelhub, "old_model_queue_threshold", 80)
)
queue_threshold = max(1, int(configured_queue_threshold or 80))
recent_model_days = max(1, int(getattr(modelhub, "recent_model_days", 7) or 7))
tasks, listing_errors = collect_active_tasks(clients, read_concurrency=read_concurrency)
log(
f"[queue-cleanup] active_scanned={len(tasks)} accounts={len(clients)} "
@@ -232,27 +334,104 @@ def cleanup_certain_oom_tasks(
read_concurrency=read_concurrency,
log=log,
)
decisions, skipped = find_certain_oom_tasks(
oom_decisions, skipped = find_certain_oom_tasks(
tasks,
repository_sizes=repository_sizes,
gpu_memory_gib=gpu_memory_gib,
)
log(
f"[queue-cleanup] certain_oom={len(decisions)} "
f"[queue-cleanup] certain_oom={len(oom_decisions)} "
f"fits={skipped['fitsKnownCapacity']} size_unknown={skipped['repositorySizeUnknown']} "
f"gpu_unknown={skipped['gpuCapacityUnknown']} dry_run={str(bool(dry_run)).lower()}"
)
oom_task_keys = {
(int(decision["accountIndex"]) - 1, int(decision["taskId"]))
for decision in oom_decisions
}
# OOM tasks are stopped first. Rank the age-policy queue as it will look
# after those certain failures are gone, so an old task moving into the
# protected first N positions is not over-cancelled.
age_rank_tasks = [
task for task in tasks if (task.account_index, task.task_id) not in oom_task_keys
]
overflow_model_ids = {
task.model_id
for account_index in range(len(clients))
if account_index not in listing_errors
for task in sorted(
(item for item in age_rank_tasks if item.account_index == account_index),
key=lambda item: item.task_id,
)[queue_threshold:]
}
model_last_modified, age_errors = _load_model_last_modified(
overflow_model_ids,
discovery=discovery,
read_concurrency=read_concurrency,
log=log,
)
old_overflow_decisions, age_skipped = find_old_overflow_tasks(
age_rank_tasks,
model_last_modified=model_last_modified,
queue_threshold=queue_threshold,
recent_model_days=recent_model_days,
reference_time=reference_time,
incomplete_accounts=set(listing_errors),
)
log(
f"[queue-cleanup] old_overflow={len(old_overflow_decisions)} "
f"threshold={queue_threshold} recent_days={recent_model_days} "
f"recent_overflow={age_skipped['recentOverflowTasks']} "
f"age_unknown={age_skipped['modelAgeUnknown']}"
)
decisions_by_key: dict[tuple[int, int], dict[str, Any]] = {}
for decision in oom_decisions:
enriched = dict(decision)
enriched["cleanupReasons"] = [decision["reason"]]
decisions_by_key[(int(decision["accountIndex"]), int(decision["taskId"]))] = enriched
for decision in old_overflow_decisions:
key = (int(decision["accountIndex"]), int(decision["taskId"]))
existing = decisions_by_key.get(key)
if existing is None:
enriched = dict(decision)
enriched["cleanupReasons"] = [decision["reason"]]
decisions_by_key[key] = enriched
continue
existing["cleanupReasons"].append(decision["reason"])
for field in ("queuePosition", "queueThreshold", "modelLastModified", "recentCutoff"):
existing[field] = decision[field]
decisions = sorted(
decisions_by_key.values(),
key=lambda item: (int(item["accountIndex"]), int(item["taskId"])),
)
cancelled: list[dict[str, Any]] = []
disappeared: list[dict[str, Any]] = []
policy_no_longer_applies: list[dict[str, Any]] = []
stop_errors: list[dict[str, Any]] = []
if decisions and not dry_run:
# Re-read each account immediately before mutation. If any active-state
# query fails for that account, fail closed and do not terminate its tasks.
refreshed_tasks, refresh_errors = collect_active_tasks(clients, read_concurrency=read_concurrency)
active_ids_by_account: dict[int, set[int]] = {}
active_positions_by_account: dict[int, dict[int, int]] = {}
for task in refreshed_tasks:
active_ids_by_account.setdefault(task.account_index, set()).add(task.task_id)
for account_index in range(len(clients)):
planned_oom_ids = {
int(decision["taskId"])
for decision in oom_decisions
if int(decision["accountIndex"]) - 1 == account_index
}
ordered_ids = sorted(
task.task_id
for task in refreshed_tasks
if task.account_index == account_index and task.task_id not in planned_oom_ids
)
active_positions_by_account[account_index] = {
task_id: position for position, task_id in enumerate(ordered_ids, start=1)
}
for account_index, details in sorted(refresh_errors.items()):
stop_errors.append(
@@ -271,6 +450,16 @@ def cleanup_certain_oom_tasks(
if int(decision["taskId"]) not in active_ids_by_account.get(account_index, set()):
disappeared.append(decision)
continue
cleanup_reasons = set(decision.get("cleanupReasons") or [decision.get("reason")])
age_only = cleanup_reasons == {"old_model_beyond_account_queue_threshold"}
current_position = active_positions_by_account.get(account_index, {}).get(int(decision["taskId"]))
if age_only and (current_position is None or current_position <= queue_threshold):
policy_no_longer_applies.append(
{**decision, "recheckedQueuePosition": current_position}
)
continue
if current_position is not None:
decision["recheckedQueuePosition"] = current_position
by_account.setdefault(account_index, []).append(decision)
batch_size = max(1, min(100, int(stop_batch_size)))
@@ -279,25 +468,36 @@ def cleanup_certain_oom_tasks(
if stop_failed:
break
decisions_by_id = {int(item["taskId"]): item for item in by_account[account_index]}
task_ids = sorted(decisions_by_id)
for batch in _chunks(task_ids, batch_size):
try:
clients[account_index].stop_tasks(batch)
except Exception as exc:
stop_errors.append(
{
"accountIndex": account_index + 1,
"taskIds": batch,
"error": f"{type(exc).__name__}: {exc}",
}
)
stop_failed = True
break
cancelled.extend(decisions_by_id[task_id] for task_id in batch)
log(
f"[queue-cleanup] account={account_index + 1:02d} "
f"cancelled_batch={len(batch)} cancelled_total={len(cancelled)}"
phase_ids = [
sorted(
task_id
for task_id, decision in decisions_by_id.items()
if ("certain_oom_repository_size_exceeds_gpu_capacity" in decision["cleanupReasons"])
== is_oom_phase
)
for is_oom_phase in (True, False)
]
for task_ids in phase_ids:
for batch in _chunks(task_ids, batch_size):
try:
clients[account_index].stop_tasks(batch)
except Exception as exc:
stop_errors.append(
{
"accountIndex": account_index + 1,
"taskIds": batch,
"error": f"{type(exc).__name__}: {exc}",
}
)
stop_failed = True
break
cancelled.extend(decisions_by_id[task_id] for task_id in batch)
log(
f"[queue-cleanup] account={account_index + 1:02d} "
f"cancelled_batch={len(batch)} cancelled_total={len(cancelled)}"
)
if stop_failed:
break
if cancelled and hasattr(modelhub, "refresh_active_counts"):
modelhub.refresh_active_counts()
@@ -309,19 +509,29 @@ def cleanup_certain_oom_tasks(
"uniqueModels": len(model_ids),
"repositorySizesComplete": len(repository_sizes),
"repositorySizeErrors": size_errors,
"modelAgeMetadataComplete": len(model_last_modified),
"modelAgeErrors": age_errors,
"listingErrors": {str(index + 1): values for index, values in listing_errors.items()},
"certainOomCount": len(decisions),
"certainOomTasks": decisions,
"certainOomCount": len(oom_decisions),
"certainOomTasks": oom_decisions,
"oldOverflowCount": len(old_overflow_decisions),
"oldOverflowTasks": old_overflow_decisions,
"oldModelQueueThreshold": queue_threshold,
"recentModelDays": recent_model_days,
"agePolicySkipped": age_skipped,
"cleanupCandidateCount": len(decisions),
"skipped": skipped,
"cancelledCount": len(cancelled),
"cancelledTasks": cancelled,
"noLongerActiveCount": len(disappeared),
"policyNoLongerAppliesCount": len(policy_no_longer_applies),
"policyNoLongerAppliesTasks": policy_no_longer_applies,
"stopErrors": stop_errors,
}
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Safely stop active ModelHub tasks that are certain to run out of GPU memory.")
parser = argparse.ArgumentParser(description="Safely stop deterministic OOM and over-threshold old ModelHub tasks.")
parser.add_argument("--execute", action="store_true", help="Actually terminate selected tasks; otherwise only print a preview")
parser.add_argument("--read-concurrency", type=int, default=6)
parser.add_argument("--stop-batch-size", type=int, default=DEFAULT_STOP_BATCH_SIZE)
@@ -353,6 +563,7 @@ def main(argv: list[str] | None = None) -> int:
write_json(report_path, summary)
print(
f"[queue-cleanup] finished certain_oom={summary['certainOomCount']} "
f"old_overflow={summary['oldOverflowCount']} "
f"cancelled={summary['cancelledCount']} stop_errors={len(summary['stopErrors'])} "
f"report={report_path}",
flush=True,