feat: dynamically clean incompatible architectures

This commit is contained in:
CoolBoy
2026-08-12 08:19:53 +08:00
parent 615bcad124
commit 7ec875563e
12 changed files with 979 additions and 132 deletions

View File

@@ -7,6 +7,7 @@ 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 hf_discovery import HuggingFaceDiscovery, inspect_repo_tree
@@ -195,6 +196,43 @@ def _load_model_last_modified(
return values, errors
def _load_model_configs(
model_ids: set[str],
*,
discovery: HuggingFaceDiscovery,
read_concurrency: int,
log: Callable[[str], None],
) -> tuple[dict[str, dict[str, Any]], dict[str, str]]:
configs: dict[str, dict[str, Any]] = {}
errors: dict[str, str] = {}
if not model_ids:
return configs, errors
completed = 0
workers = min(max(1, int(read_concurrency)), len(model_ids))
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = {
executor.submit(discovery.get_model_config, model_id): model_id
for model_id in sorted(model_ids)
}
for future in as_completed(futures):
model_id = futures[future]
try:
config, fetch_error = future.result()
if fetch_error or not isinstance(config, dict) or not config:
errors[model_id] = str(fetch_error or "model_config_empty")
else:
configs[model_id] = dict(config)
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] architecture_scan={completed}/{len(model_ids)} "
f"complete={len(configs)} unknown={len(errors)}"
)
return configs, errors
def find_certain_oom_tasks(
tasks: list[OwnedTask],
*,
@@ -238,6 +276,94 @@ def find_certain_oom_tasks(
return decisions, skipped
def find_architecture_incompatible_tasks(
tasks: list[OwnedTask],
*,
architecture_blocks: dict[str, dict[str, Any]],
task_contexts: dict[str, dict[str, Any]],
model_configs: dict[str, dict[str, Any]],
) -> tuple[list[dict[str, Any]], dict[str, int]]:
"""Select waiting tasks that exactly match a learned compatibility block."""
decisions: list[dict[str, Any]] = []
skipped = {
"submissionContextUnknown": 0,
"submissionContextMismatch": 0,
"modelArchitectureUnknown": 0,
"noMatchingBlock": 0,
"runningMatchedProtected": 0,
}
if not architecture_blocks:
return decisions, skipped
for task in tasks:
context = task_contexts.get(str(task.task_id))
if not isinstance(context, dict):
skipped["submissionContextUnknown"] += 1
continue
context_model = str(context.get("modelId") or "").strip()
context_gpu = str(context.get("targetGpu") or "").strip()
framework = str(context.get("framework") or "").strip()
task_type = str(context.get("taskType") or "").strip()
if (
not framework
or not task_type
or (context_model and context_model != task.model_id)
or (context_gpu and context_gpu.casefold() != task.gpu_type.casefold())
):
skipped["submissionContextMismatch"] += 1
continue
profile_data = context.get("modelProfile")
if not isinstance(profile_data, dict):
profile_data = {}
model_type = profile_data.get("modelType")
architectures = profile_data.get("architectures")
if not model_type and not architectures:
config = model_configs.get(task.model_id) or {}
model_type = config.get("model_type")
architectures = config.get("architectures")
profiles = architecture_profiles(model_type, architectures)
if not profiles:
skipped["modelArchitectureUnknown"] += 1
continue
matching_block: dict[str, Any] | None = None
for profile in profiles:
key = architecture_compatibility_key(
task.gpu_type,
framework,
task_type,
profile["signature"],
)
block = architecture_blocks.get(key or "")
if isinstance(block, dict):
matching_block = block
break
if matching_block is None:
skipped["noMatchingBlock"] += 1
continue
if task.status != "waiting":
skipped["runningMatchedProtected"] += 1
continue
decisions.append(
{
"accountIndex": task.account_index + 1,
"taskId": task.task_id,
"modelId": task.model_id,
"gpuType": task.gpu_type,
"framework": framework,
"taskType": task_type,
"status": task.status,
"architectureSignature": matching_block.get("architectureSignature"),
"architectureMatchType": matching_block.get("matchType"),
"architectureBlockExpiresAt": matching_block.get("expiresAt"),
"architectureBlockEvidenceCount": matching_block.get("evidenceCount"),
"reason": "known_framework_architecture_incompatible",
}
)
return decisions, skipped
def find_old_overflow_tasks(
tasks: list[OwnedTask],
*,
@@ -319,11 +445,14 @@ 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,
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 tasks and old tasks beyond each account's protected prefix."""
"""Stop deterministic OOM/architecture tasks and over-threshold old tasks."""
clients = list(modelhub.clients)
reference_time = reference_time or utc_now()
configured_reserved_slots = (
@@ -359,32 +488,97 @@ def cleanup_certain_oom_tasks(
}
model_ids = {task.model_id for task in tasks}
repository_sizes, size_errors = _load_repository_sizes(
model_ids,
discovery=discovery,
read_concurrency=read_concurrency,
log=log,
)
oom_decisions, skipped = find_certain_oom_tasks(
tasks,
repository_sizes=repository_sizes,
gpu_memory_gib=gpu_memory_gib,
)
repository_sizes: dict[str, int] = {}
size_errors: dict[str, str] = {}
oom_decisions: list[dict[str, Any]] = []
skipped = {
"repositorySizeUnknown": 0,
"gpuCapacityUnknown": 0,
"fitsKnownCapacity": 0,
}
if not architecture_only:
repository_sizes, size_errors = _load_repository_sizes(
model_ids,
discovery=discovery,
read_concurrency=read_concurrency,
log=log,
)
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(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
architecture_blocks = (
architecture_compatibility_blocks
if isinstance(architecture_compatibility_blocks, dict)
else {}
)
task_contexts = (
task_compatibility_contexts
if isinstance(task_compatibility_contexts, dict)
else {}
)
block_combinations = {
(
str(block.get("targetGpu") or "").strip().casefold(),
str(block.get("framework") or "").strip().casefold(),
str(block.get("taskType") or "").strip().casefold(),
)
for block in architecture_blocks.values()
if isinstance(block, dict)
}
# 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.
architecture_model_ids: set[str] = set()
for task in tasks:
context = task_contexts.get(str(task.task_id))
if not isinstance(context, dict):
continue
combination = (
task.gpu_type.casefold(),
str(context.get("framework") or "").strip().casefold(),
str(context.get("taskType") or "").strip().casefold(),
)
profile = context.get("modelProfile")
profile = profile if isinstance(profile, dict) else {}
if combination in block_combinations and not (
profile.get("modelType") or profile.get("architectures")
):
architecture_model_ids.add(task.model_id)
model_configs, model_config_errors = _load_model_configs(
architecture_model_ids,
discovery=discovery,
read_concurrency=read_concurrency,
log=log,
)
architecture_decisions, architecture_skipped = find_architecture_incompatible_tasks(
tasks,
architecture_blocks=architecture_blocks,
task_contexts=task_contexts,
model_configs=model_configs,
)
log(
f"[queue-cleanup] architecture_incompatible={len(architecture_decisions)} "
f"blocks={len(architecture_blocks)} "
f"context_unknown={architecture_skipped['submissionContextUnknown']} "
f"architecture_unknown={architecture_skipped['modelArchitectureUnknown']} "
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 oom_task_keys
task
for task in tasks
if (task.account_index, task.task_id) not in deterministic_task_keys
]
overflow_model_ids = {
task.model_id
@@ -396,20 +590,32 @@ def cleanup_certain_oom_tasks(
)[queue_thresholds.get(account_index, len(age_rank_tasks)):]
if task.status == "waiting"
}
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),
)
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))} "
@@ -420,10 +626,22 @@ def cleanup_certain_oom_tasks(
)
decisions_by_key: dict[tuple[int, int], dict[str, Any]] = {}
for decision in oom_decisions:
for decision in [*oom_decisions, *architecture_decisions]:
enriched = dict(decision)
enriched["cleanupReasons"] = [decision["reason"]]
decisions_by_key[(int(decision["accountIndex"]), int(decision["taskId"]))] = enriched
key = (int(decision["accountIndex"]), int(decision["taskId"]))
existing = decisions_by_key.get(key)
if existing is None:
decisions_by_key[key] = enriched
continue
existing["cleanupReasons"].append(decision["reason"])
existing.update(
{
field: value
for field, value in decision.items()
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)
@@ -455,15 +673,16 @@ def cleanup_certain_oom_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_oom_ids = {
planned_deterministic_ids = {
int(decision["taskId"])
for decision in oom_decisions
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_oom_ids
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)
@@ -488,6 +707,10 @@ def cleanup_certain_oom_tasks(
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)
@@ -510,6 +733,16 @@ def cleanup_certain_oom_tasks(
}
)
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)
@@ -520,17 +753,36 @@ def cleanup_certain_oom_tasks(
if stop_failed:
break
decisions_by_id = {int(item["taskId"]): item for item in by_account[account_index]}
for is_oom_phase in (True, False):
for cleanup_phase in ("oom", "architecture", "age"):
task_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
if (
cleanup_phase == "oom"
and "certain_oom_repository_size_exceeds_gpu_capacity"
in decision["cleanupReasons"]
)
or (
cleanup_phase == "architecture"
and "certain_oom_repository_size_exceeds_gpu_capacity"
not in decision["cleanupReasons"]
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 not is_oom_phase and task_ids:
# OOM stops can change actual positions, and a waiting task
# can start running after the account-wide recheck above.
# Re-read this account immediately before its age-only stop.
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.
try:
phase_tasks: dict[int, OwnedTask] = {}
for status in ACTIVE_FILTER_STATUSES:
@@ -545,7 +797,7 @@ def cleanup_certain_oom_tasks(
{
"accountIndex": account_index + 1,
"taskIds": task_ids,
"error": f"age_policy_final_recheck_failed: {type(exc).__name__}: {exc}",
"error": f"policy_final_recheck_failed: {type(exc).__name__}: {exc}",
}
)
stop_failed = True
@@ -563,12 +815,24 @@ def cleanup_certain_oom_tasks(
disappeared.append(decisions_by_id[task_id])
continue
current_position = phase_positions.get(task_id)
if (
account_queue_threshold is None
or current_position is None
or current_position <= account_queue_threshold
or current_task.status != "waiting"
):
reasons = set(
decisions_by_id[task_id].get("cleanupReasons")
or [decisions_by_id[task_id].get("reason")]
)
architecture_applies = bool(
cleanup_phase == "architecture"
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:
policy_no_longer_applies.append(
{
**decisions_by_id[task_id],
@@ -577,7 +841,7 @@ def cleanup_certain_oom_tasks(
"policyChangeReason": (
"task_started_running"
if current_task.status == "running"
else "queue_position_or_status_changed"
else "queue_position_status_or_policy_changed"
),
}
)
@@ -612,6 +876,7 @@ def cleanup_certain_oom_tasks(
return {
"dryRun": bool(dry_run),
"architectureOnly": bool(architecture_only),
"accounts": len(clients),
"activeScanned": len(tasks),
"uniqueModels": len(model_ids),
@@ -622,6 +887,12 @@ def cleanup_certain_oom_tasks(
"listingErrors": {str(index + 1): values for index, values in listing_errors.items()},
"certainOomCount": len(oom_decisions),
"certainOomTasks": oom_decisions,
"architectureBlockCount": len(architecture_blocks),
"architectureIncompatibleCount": len(architecture_decisions),
"architectureIncompatibleTasks": architecture_decisions,
"architectureModelConfigsComplete": len(model_configs),
"architectureModelConfigErrors": model_config_errors,
"architecturePolicySkipped": architecture_skipped,
"oldOverflowCount": len(old_overflow_decisions),
"oldOverflowTasks": old_overflow_decisions,
"oldModelQueueThresholds": [
@@ -646,7 +917,9 @@ def cleanup_certain_oom_tasks(
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Safely stop deterministic OOM and over-threshold old ModelHub tasks.")
parser = argparse.ArgumentParser(
description="Safely stop deterministic OOM/architecture 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)
@@ -678,6 +951,7 @@ def main(argv: list[str] | None = None) -> int:
write_json(report_path, summary)
print(
f"[queue-cleanup] finished certain_oom={summary['certainOomCount']} "
f"architecture_incompatible={summary['architectureIncompatibleCount']} "
f"old_overflow={summary['oldOverflowCount']} "
f"cancelled={summary['cancelledCount']} stop_errors={len(summary['stopErrors'])} "
f"report={report_path}",