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

feat: make model age admission-only

Browse Source
This commit is contained in:
CoolBoy
2026-08-12 10:18:09 +08:00
parent 86088ce577
commit 91d1d3d87d
15 changed files with 438 additions and 653 deletions
Download Patch File Download Diff File Expand all files Collapse all files

284
modelhub_submmit_api/queue_cleanup.py
View File

@@ -3,13 +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 architecture_compatibility import architecture_compatibility_key, architecture_profiles
from candidate_preflight import CandidatePreflightAdvisor, MODEL_LOAD_OVERHEAD
from common import utc_now, write_json
from common import 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
@@ -168,41 +167,6 @@ 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 _load_model_configs(
model_ids: set[str],
*,
@@ -480,74 +444,6 @@ def find_architecture_incompatible_tasks(
return decisions, skipped
def find_old_overflow_tasks(
tasks: list[OwnedTask],
*,
model_last_modified: dict[str, datetime],
queue_threshold: int | dict[int, 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]]:
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,
"accountThresholdUnknown": 0,
"runningOverflowProtected": 0,
}
for account_index, account_tasks in sorted(grouped.items()):
if account_index in incomplete_accounts:
continue
if isinstance(queue_threshold, dict):
configured_threshold = queue_threshold.get(account_index)
if configured_threshold is None:
skipped["accountThresholdUnknown"] += 1
continue
else:
configured_threshold = queue_threshold
threshold = max(0, int(configured_threshold))
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
if task.status != "waiting":
skipped["runningOverflowProtected"] += 1
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]
@@ -564,18 +460,11 @@ def cleanup_certain_oom_tasks(
architecture_compatibility_blocks: dict[str, dict[str, Any]] | None = None,
task_compatibility_contexts: dict[str, dict[str, Any]] | None = None,
architecture_only: bool = False,
age_reserved_slots: int | None = None,
reference_time: datetime | None = None,
log: Callable[[str], None] = print,
) -> dict[str, Any]:
"""Stop deterministic OOM/architecture tasks and over-threshold old tasks."""
"""Stop only deterministically impossible OOM/architecture tasks."""
clients = list(modelhub.clients)
reference_time = reference_time or utc_now()
configured_reserved_slots = (
age_reserved_slots
if age_reserved_slots is not None
else getattr(modelhub, "recent_model_reserve_slots", 10)
)
configured_reserved_slots = getattr(modelhub, "recent_model_reserve_slots", 5)
reserved_slots = max(0, int(configured_reserved_slots or 0))
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)
@@ -604,6 +493,11 @@ def cleanup_certain_oom_tasks(
index: max(0, int(account_caps[index]) - reserved_slots)
for index in range(min(len(clients), len(account_caps)))
}
log(
f"[queue-cleanup] age_cleanup=disabled deterministic_cleanup=enabled "
f"age_policy=admission_only reserve_recent_slots={reserved_slots} "
f"recent_days={recent_model_days}"
)
model_ids = {task.model_id for task in tasks}
repository_sizes: dict[str, int] = {}
@@ -717,62 +611,6 @@ def cleanup_certain_oom_tasks(
f"running_protected={architecture_skipped['runningMatchedProtected']}"
)
deterministic_task_keys = {
(int(decision["accountIndex"]) - 1, int(decision["taskId"]))
for decision in [*oom_decisions, *architecture_decisions]
}
# Deterministically impossible tasks are stopped first. Rank the age-policy
# queue as it will look afterwards to avoid over-cancelling old models.
age_rank_tasks = [
task
for task in tasks
if (task.account_index, task.task_id) not in deterministic_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_thresholds.get(account_index, len(age_rank_tasks)):]
if task.status == "waiting"
}
model_last_modified: dict[str, datetime] = {}
age_errors: dict[str, str] = {}
old_overflow_decisions: list[dict[str, Any]] = []
age_skipped = {
"accountsWithIncompleteListing": len(listing_errors),
"withinFirstQueuePositions": 0,
"recentOverflowTasks": 0,
"modelAgeUnknown": 0,
"accountThresholdUnknown": 0,
"runningOverflowProtected": 0,
}
if not architecture_only:
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_thresholds,
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"thresholds={','.join(str(queue_thresholds[index]) for index in sorted(queue_thresholds))} "
f"reserve_recent_slots={reserved_slots} recent_days={recent_model_days} "
f"recent_overflow={age_skipped['recentOverflowTasks']} "
f"age_unknown={age_skipped['modelAgeUnknown']} "
f"running_protected={age_skipped['runningOverflowProtected']}"
)
decisions_by_key: dict[tuple[int, int], dict[str, Any]] = {}
for decision in [*oom_decisions, *architecture_decisions]:
enriched = dict(decision)
@@ -790,17 +628,6 @@ def cleanup_certain_oom_tasks(
if field not in {"reason", "cleanupReasons"} and value is not None
}
)
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"])),
@@ -815,26 +642,10 @@ def cleanup_certain_oom_tasks(
# 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]] = {}
active_status_by_account: dict[int, dict[int, str]] = {}
for task in refreshed_tasks:
active_ids_by_account.setdefault(task.account_index, set()).add(task.task_id)
active_status_by_account.setdefault(task.account_index, {})[task.task_id] = task.status
for account_index in range(len(clients)):
planned_deterministic_ids = {
int(decision["taskId"])
for decision in [*oom_decisions, *architecture_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_deterministic_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(
@@ -854,45 +665,20 @@ def cleanup_certain_oom_tasks(
disappeared.append(decision)
continue
cleanup_reasons = set(decision.get("cleanupReasons") or [decision.get("reason")])
age_only = cleanup_reasons == {"old_model_beyond_account_queue_threshold"}
architecture_without_oom = bool(
"known_framework_architecture_incompatible" in cleanup_reasons
and "certain_oom_repository_size_exceeds_gpu_capacity" not in cleanup_reasons
)
current_position = active_positions_by_account.get(account_index, {}).get(int(decision["taskId"]))
current_status = active_status_by_account.get(account_index, {}).get(int(decision["taskId"]))
account_queue_threshold = queue_thresholds.get(account_index)
if age_only and (
current_position is None
or account_queue_threshold is None
or current_position <= account_queue_threshold
or current_status != "waiting"
):
policy_no_longer_applies.append(
{
**decision,
"recheckedQueuePosition": current_position,
"recheckedStatus": current_status,
"policyChangeReason": (
"task_started_running"
if current_status == "running"
else "queue_position_or_status_changed"
),
}
)
continue
if architecture_without_oom and current_status != "waiting":
policy_no_longer_applies.append(
{
**decision,
"recheckedQueuePosition": current_position,
"recheckedStatus": current_status,
"policyChangeReason": "task_started_running",
}
)
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)))
@@ -901,7 +687,7 @@ def cleanup_certain_oom_tasks(
if stop_failed:
break
decisions_by_id = {int(item["taskId"]): item for item in by_account[account_index]}
for cleanup_phase in ("oom", "architecture", "age"):
for cleanup_phase in ("oom", "architecture"):
task_ids = sorted(
task_id
for task_id, decision in decisions_by_id.items()
@@ -917,20 +703,10 @@ def cleanup_certain_oom_tasks(
and "known_framework_architecture_incompatible"
in decision["cleanupReasons"]
)
or (
cleanup_phase == "age"
and "certain_oom_repository_size_exceeds_gpu_capacity"
not in decision["cleanupReasons"]
and "known_framework_architecture_incompatible"
not in decision["cleanupReasons"]
and "old_model_beyond_account_queue_threshold"
in decision["cleanupReasons"]
)
)
if cleanup_phase != "oom" and task_ids:
# Earlier deterministic stops can change positions, and a
# waiting architecture/age task can start running after the
# account-wide recheck. Re-read before each later phase.
# A waiting architecture task can start running after the
# account-wide recheck. Re-read before this protected phase.
try:
phase_tasks: dict[int, OwnedTask] = {}
for status in ACTIVE_FILTER_STATUSES:
@@ -951,18 +727,12 @@ def cleanup_certain_oom_tasks(
stop_failed = True
break
phase_positions = {
task_id: position
for position, task_id in enumerate(sorted(phase_tasks), start=1)
}
eligible_task_ids: list[int] = []
account_queue_threshold = queue_thresholds.get(account_index)
for task_id in task_ids:
current_task = phase_tasks.get(task_id)
if current_task is None:
disappeared.append(decisions_by_id[task_id])
continue
current_position = phase_positions.get(task_id)
reasons = set(
decisions_by_id[task_id].get("cleanupReasons")
or [decisions_by_id[task_id].get("reason")]
@@ -972,29 +742,19 @@ def cleanup_certain_oom_tasks(
and "known_framework_architecture_incompatible" in reasons
and current_task.status == "waiting"
)
age_applies = bool(
cleanup_phase == "age"
and "old_model_beyond_account_queue_threshold" in reasons
and account_queue_threshold is not None
and current_position is not None
and current_position > account_queue_threshold
and current_task.status == "waiting"
)
if not architecture_applies and not age_applies:
if not architecture_applies:
policy_no_longer_applies.append(
{
**decisions_by_id[task_id],
"recheckedQueuePosition": current_position,
"recheckedStatus": current_task.status,
"policyChangeReason": (
"task_started_running"
if current_task.status == "running"
else "queue_position_status_or_policy_changed"
else "task_no_longer_architecture_cleanup_eligible"
),
}
)
continue
decisions_by_id[task_id]["recheckedQueuePosition"] = current_position
eligible_task_ids.append(task_id)
task_ids = eligible_task_ids
@@ -1030,8 +790,9 @@ def cleanup_certain_oom_tasks(
"uniqueModels": len(model_ids),
"repositorySizesComplete": len(repository_sizes),
"repositorySizeErrors": size_errors,
"modelAgeMetadataComplete": len(model_last_modified),
"modelAgeErrors": age_errors,
"ageCleanupMode": "admission_only",
"modelAgeMetadataComplete": 0,
"modelAgeErrors": {},
"listingErrors": {str(index + 1): values for index, values in listing_errors.items()},
"certainOomCount": len(oom_decisions),
"certainOomTasks": oom_decisions,
@@ -1046,8 +807,8 @@ def cleanup_certain_oom_tasks(
},
"architectureFrameworkCatalogErrors": framework_catalog_errors,
"architecturePolicySkipped": architecture_skipped,
"oldOverflowCount": len(old_overflow_decisions),
"oldOverflowTasks": old_overflow_decisions,
"oldOverflowCount": 0,
"oldOverflowTasks": [],
"oldModelQueueThresholds": [
queue_thresholds.get(index) for index in range(len(clients))
],
@@ -1057,7 +818,10 @@ def cleanup_certain_oom_tasks(
],
"recentModelReserveSlots": reserved_slots,
"recentModelDays": recent_model_days,
"agePolicySkipped": age_skipped,
"agePolicySkipped": {
"cleanupDisabled": True,
"reason": "admission_only",
},
"cleanupCandidateCount": len(decisions),
"skipped": skipped,
"cancelledCount": len(cancelled),
@@ -1071,7 +835,7 @@ def cleanup_certain_oom_tasks(
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="Safely stop deterministic OOM/architecture and over-threshold old ModelHub tasks."
description="Safely stop only deterministic OOM/architecture 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)
@@ -1105,7 +869,7 @@ def main(argv: list[str] | None = None) -> int:
print(
f"[queue-cleanup] finished certain_oom={summary['certainOomCount']} "
f"architecture_incompatible={summary['architectureIncompatibleCount']} "
f"old_overflow={summary['oldOverflowCount']} "
f"age_cleanup={summary['ageCleanupMode']} "
f"cancelled={summary['cancelledCount']} stop_errors={len(summary['stopErrors'])} "
f"report={report_path}",
flush=True,