1118 lines
46 KiB
Python
1118 lines
46 KiB
Python
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 hf_discovery import HuggingFaceDiscovery, inspect_repo_tree
|
|
from modelhub_client import ModelHubClient, ModelHubClientPool
|
|
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
|
|
from task_registry import task_type_from_history_task
|
|
|
|
|
|
ACTIVE_FILTER_STATUSES = ("waiting", "running")
|
|
DEFAULT_STOP_BATCH_SIZE = 50
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OwnedTask:
|
|
account_index: int
|
|
task_id: int
|
|
model_id: str
|
|
gpu_type: str
|
|
status: str
|
|
task_type: str = ""
|
|
framework: str = ""
|
|
|
|
|
|
def _normalize_status(value: Any) -> str:
|
|
return str(value or "").strip().lower()
|
|
|
|
|
|
def _parse_owned_task(account_index: int, record: dict[str, Any]) -> OwnedTask | None:
|
|
task_id = record.get("taskId")
|
|
model_id = str(record.get("modelId") or "").strip()
|
|
gpu_type = str(record.get("gpuType") or "").strip()
|
|
status = _normalize_status(record.get("status"))
|
|
if not model_id or not gpu_type or status not in ACTIVE_FILTER_STATUSES:
|
|
return None
|
|
try:
|
|
numeric_task_id = int(str(task_id).strip())
|
|
except (TypeError, ValueError):
|
|
return None
|
|
if numeric_task_id <= 0:
|
|
return None
|
|
return OwnedTask(
|
|
account_index=account_index,
|
|
task_id=numeric_task_id,
|
|
model_id=model_id,
|
|
gpu_type=gpu_type,
|
|
status=status,
|
|
task_type=task_type_from_history_task(record) or "",
|
|
# The current API does not expose this field today, but retaining it
|
|
# makes the cleanup automatically use it if the platform adds it.
|
|
framework=str(record.get("framework") or "").strip(),
|
|
)
|
|
|
|
|
|
def _fetch_status_tasks(
|
|
client: ModelHubClient,
|
|
*,
|
|
account_index: int,
|
|
status: str,
|
|
page_size: int = 100,
|
|
) -> list[OwnedTask]:
|
|
current = 1
|
|
results: list[OwnedTask] = []
|
|
while True:
|
|
payload = client.list_tasks_page(
|
|
current=current,
|
|
page_size=page_size,
|
|
only_mine=True,
|
|
status=status,
|
|
)
|
|
page = payload.get("data") or {}
|
|
records = page.get("records") or []
|
|
if not isinstance(records, list):
|
|
raise ValueError("ModelHub task page records are invalid")
|
|
for record in records:
|
|
if not isinstance(record, dict):
|
|
continue
|
|
task = _parse_owned_task(account_index, record)
|
|
# The platform has silently ignored unknown status filters before.
|
|
# Accept only records whose returned state is explicitly active.
|
|
if task is not None:
|
|
results.append(task)
|
|
pages = int(page.get("pages") or 0)
|
|
if not records or current >= pages:
|
|
break
|
|
current += 1
|
|
return results
|
|
|
|
|
|
def collect_active_tasks(
|
|
clients: list[ModelHubClient],
|
|
*,
|
|
read_concurrency: int = 6,
|
|
statuses: Iterable[str] = ACTIVE_FILTER_STATUSES,
|
|
) -> tuple[list[OwnedTask], dict[int, list[str]]]:
|
|
"""Fetch active tasks from every account without sharing account-scoped reads."""
|
|
requested_statuses = tuple(dict.fromkeys(_normalize_status(status) for status in statuses))
|
|
tasks: list[OwnedTask] = []
|
|
errors: dict[int, list[str]] = {}
|
|
jobs = [(index, client, status) for index, client in enumerate(clients) for status in requested_statuses]
|
|
workers = min(max(1, int(read_concurrency)), max(1, len(jobs)))
|
|
with ThreadPoolExecutor(max_workers=workers) as executor:
|
|
futures = {
|
|
executor.submit(
|
|
_fetch_status_tasks,
|
|
client,
|
|
account_index=index,
|
|
status=status,
|
|
): (index, status)
|
|
for index, client, status in jobs
|
|
}
|
|
for future in as_completed(futures):
|
|
index, status = futures[future]
|
|
try:
|
|
tasks.extend(future.result())
|
|
except Exception as exc:
|
|
errors.setdefault(index, []).append(f"{status}: {type(exc).__name__}: {exc}")
|
|
|
|
deduped: dict[tuple[int, int], OwnedTask] = {}
|
|
for task in tasks:
|
|
deduped[(task.account_index, task.task_id)] = task
|
|
return sorted(deduped.values(), key=lambda item: (item.account_index, item.task_id)), errors
|
|
|
|
|
|
def _load_repository_sizes(
|
|
model_ids: set[str],
|
|
*,
|
|
discovery: HuggingFaceDiscovery,
|
|
read_concurrency: int,
|
|
log: Callable[[str], None],
|
|
) -> tuple[dict[str, int], dict[str, str]]:
|
|
sizes: dict[str, int] = {}
|
|
errors: dict[str, str] = {}
|
|
completed = 0
|
|
workers = min(max(1, int(read_concurrency)), max(1, len(model_ids)))
|
|
|
|
def inspect(model_id: str) -> tuple[str, int | None]:
|
|
tree = discovery.list_repo_tree(model_id)
|
|
return model_id, inspect_repo_tree(model_id, tree).repository_size_bytes
|
|
|
|
with ThreadPoolExecutor(max_workers=workers) as executor:
|
|
futures = {executor.submit(inspect, model_id): model_id for model_id in sorted(model_ids)}
|
|
for future in as_completed(futures):
|
|
model_id = futures[future]
|
|
try:
|
|
returned_model_id, size = future.result()
|
|
if size is None:
|
|
errors[returned_model_id] = "recursive_repository_size_incomplete"
|
|
else:
|
|
sizes[returned_model_id] = size
|
|
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] size_scan={completed}/{len(model_ids)} "
|
|
f"complete={len(sizes)} unknown={len(errors)}"
|
|
)
|
|
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],
|
|
*,
|
|
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 _load_framework_catalog(
|
|
combinations: set[tuple[str, str]],
|
|
*,
|
|
modelhub: ModelHubClientPool,
|
|
read_concurrency: int,
|
|
log: Callable[[str], None],
|
|
) -> tuple[dict[tuple[str, str], set[str]], dict[str, str]]:
|
|
"""Load the live framework set for (GPU, task type) pairs.
|
|
|
|
Keys are normalized so active records, learned blocks, and API rows can be
|
|
compared without relying on platform capitalization.
|
|
"""
|
|
catalog: dict[tuple[str, str], set[str]] = {}
|
|
errors: dict[str, str] = {}
|
|
if not combinations:
|
|
return catalog, errors
|
|
|
|
targets = {
|
|
(str(gpu).strip().casefold(), str(task_type).strip().casefold()): (
|
|
str(gpu).strip(),
|
|
str(task_type).strip(),
|
|
)
|
|
for gpu, task_type in combinations
|
|
if str(gpu).strip() and str(task_type).strip()
|
|
}
|
|
workers = min(max(1, int(read_concurrency)), max(1, len(targets)))
|
|
|
|
def fetch(gpu: str, task_type: str) -> tuple[tuple[str, str], set[str]]:
|
|
rows = modelhub.list_framework_stats(task_type, gpu)
|
|
frameworks = {
|
|
str(row.get("framework") or "").strip()
|
|
for row in rows
|
|
if isinstance(row, dict) and str(row.get("framework") or "").strip()
|
|
}
|
|
return (gpu.casefold(), task_type.casefold()), frameworks
|
|
|
|
with ThreadPoolExecutor(max_workers=workers) as executor:
|
|
futures = {
|
|
executor.submit(fetch, gpu, task_type): normalized_key
|
|
for normalized_key, (gpu, task_type) in sorted(targets.items())
|
|
}
|
|
for future in as_completed(futures):
|
|
gpu, task_type = futures[future]
|
|
key = f"{gpu}|{task_type}"
|
|
try:
|
|
returned_key, frameworks = future.result()
|
|
if frameworks:
|
|
catalog[returned_key] = frameworks
|
|
else:
|
|
errors[key] = "live_framework_catalog_empty"
|
|
except Exception as exc:
|
|
errors[key] = f"{type(exc).__name__}: {exc}"
|
|
|
|
log(
|
|
f"[queue-cleanup] framework_catalog={len(catalog)}/{len(targets)} "
|
|
f"frameworks={sum(len(items) for items in catalog.values())} "
|
|
f"unknown={len(errors)}"
|
|
)
|
|
return catalog, errors
|
|
|
|
|
|
def find_certain_oom_tasks(
|
|
tasks: list[OwnedTask],
|
|
*,
|
|
repository_sizes: dict[str, int],
|
|
gpu_memory_gib: dict[str, float] | None = None,
|
|
) -> tuple[list[dict[str, Any]], dict[str, int]]:
|
|
capacities = dict(CandidatePreflightAdvisor(gpu_memory_gib=gpu_memory_gib).gpu_memory_gib)
|
|
decisions: list[dict[str, Any]] = []
|
|
skipped = {
|
|
"repositorySizeUnknown": 0,
|
|
"gpuCapacityUnknown": 0,
|
|
"fitsKnownCapacity": 0,
|
|
}
|
|
for task in tasks:
|
|
size_bytes = repository_sizes.get(task.model_id)
|
|
if size_bytes is None:
|
|
skipped["repositorySizeUnknown"] += 1
|
|
continue
|
|
capacity_gib = capacities.get(task.gpu_type)
|
|
if capacity_gib is None:
|
|
skipped["gpuCapacityUnknown"] += 1
|
|
continue
|
|
repository_gib = size_bytes / (1024**3)
|
|
required_gib = repository_gib * MODEL_LOAD_OVERHEAD
|
|
if required_gib <= capacity_gib:
|
|
skipped["fitsKnownCapacity"] += 1
|
|
continue
|
|
decisions.append(
|
|
{
|
|
"accountIndex": task.account_index + 1,
|
|
"taskId": task.task_id,
|
|
"modelId": task.model_id,
|
|
"gpuType": task.gpu_type,
|
|
"status": task.status,
|
|
"repositorySizeGiB": round(repository_gib, 3),
|
|
"requiredGiB": round(required_gib, 3),
|
|
"gpuCapacityGiB": float(capacity_gib),
|
|
"reason": "certain_oom_repository_size_exceeds_gpu_capacity",
|
|
}
|
|
)
|
|
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]],
|
|
framework_catalog: dict[tuple[str, str], set[str]] | None = None,
|
|
) -> tuple[list[dict[str, Any]], dict[str, int]]:
|
|
"""Select tasks with exact or safely exhaustive learned incompatibility."""
|
|
decisions: list[dict[str, Any]] = []
|
|
skipped = {
|
|
"submissionContextUnknown": 0,
|
|
"submissionContextMismatch": 0,
|
|
"frameworkContextUnknown": 0,
|
|
"frameworkCatalogUnknown": 0,
|
|
"partiallyBlockedFrameworkSet": 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))
|
|
context = context if isinstance(context, dict) else {}
|
|
context_model = str(context.get("modelId") or "").strip()
|
|
context_gpu = str(context.get("targetGpu") or "").strip()
|
|
framework = str(context.get("framework") or task.framework or "").strip()
|
|
task_type = str(context.get("taskType") or task.task_type or "").strip()
|
|
if (
|
|
(context_model and context_model != task.model_id)
|
|
or (context_gpu and context_gpu.casefold() != task.gpu_type.casefold())
|
|
):
|
|
skipped["submissionContextMismatch"] += 1
|
|
continue
|
|
if not task_type:
|
|
skipped["submissionContextUnknown"] += 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
|
|
|
|
def matching_block_for(candidate_framework: str) -> dict[str, Any] | None:
|
|
for profile in profiles:
|
|
key = architecture_compatibility_key(
|
|
task.gpu_type,
|
|
candidate_framework,
|
|
task_type,
|
|
profile["signature"],
|
|
)
|
|
block = architecture_blocks.get(key or "")
|
|
if isinstance(block, dict):
|
|
return block
|
|
return None
|
|
|
|
matching_blocks: dict[str, dict[str, Any]] = {}
|
|
match_scope = "exact_framework"
|
|
if framework:
|
|
matching_block = matching_block_for(framework)
|
|
if matching_block is not None:
|
|
matching_blocks[framework] = matching_block
|
|
else:
|
|
skipped["frameworkContextUnknown"] += 1
|
|
available_frameworks = sorted(
|
|
(framework_catalog or {}).get(
|
|
(task.gpu_type.casefold(), task_type.casefold()),
|
|
set(),
|
|
),
|
|
key=str.casefold,
|
|
)
|
|
if not available_frameworks:
|
|
skipped["frameworkCatalogUnknown"] += 1
|
|
continue
|
|
for candidate_framework in available_frameworks:
|
|
matching_block = matching_block_for(candidate_framework)
|
|
if matching_block is not None:
|
|
matching_blocks[candidate_framework] = matching_block
|
|
if len(matching_blocks) != len(available_frameworks):
|
|
if matching_blocks:
|
|
skipped["partiallyBlockedFrameworkSet"] += 1
|
|
else:
|
|
skipped["noMatchingBlock"] += 1
|
|
continue
|
|
framework = "*"
|
|
match_scope = "all_live_frameworks"
|
|
|
|
if not matching_blocks:
|
|
skipped["noMatchingBlock"] += 1
|
|
continue
|
|
if task.status != "waiting":
|
|
skipped["runningMatchedProtected"] += 1
|
|
continue
|
|
representative_block = next(iter(matching_blocks.values()))
|
|
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,
|
|
"architectureMatchScope": match_scope,
|
|
"evaluatedFrameworks": sorted(matching_blocks, key=str.casefold),
|
|
"architectureSignatures": sorted(
|
|
{
|
|
str(block.get("architectureSignature") or "")
|
|
for block in matching_blocks.values()
|
|
if block.get("architectureSignature")
|
|
}
|
|
),
|
|
"architectureSignature": representative_block.get("architectureSignature"),
|
|
"architectureMatchType": representative_block.get("matchType"),
|
|
"architectureBlockExpiresAt": representative_block.get("expiresAt"),
|
|
"architectureBlockEvidenceCount": sum(
|
|
int(block.get("evidenceCount") or 0)
|
|
for block in matching_blocks.values()
|
|
),
|
|
"reason": "known_framework_architecture_incompatible",
|
|
}
|
|
)
|
|
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]
|
|
|
|
|
|
def cleanup_certain_oom_tasks(
|
|
modelhub: ModelHubClientPool,
|
|
discovery: HuggingFaceDiscovery,
|
|
*,
|
|
dry_run: bool = False,
|
|
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/architecture tasks and over-threshold old 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)
|
|
)
|
|
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)
|
|
log(
|
|
f"[queue-cleanup] active_scanned={len(tasks)} accounts={len(clients)} "
|
|
f"listing_errors={sum(len(items) for items in listing_errors.values())} "
|
|
f"task_type_recovered={sum(1 for task in tasks if task.task_type)} "
|
|
f"framework_exposed={sum(1 for task in tasks if task.framework)}"
|
|
)
|
|
|
|
observed_active_counts: list[int | None] = [0 for _ in clients]
|
|
for task in tasks:
|
|
current_count = observed_active_counts[task.account_index]
|
|
if current_count is not None:
|
|
observed_active_counts[task.account_index] = current_count + 1
|
|
for account_index in listing_errors:
|
|
observed_active_counts[account_index] = None
|
|
if hasattr(modelhub, "observe_capacity_lower_bounds"):
|
|
account_caps = list(modelhub.observe_capacity_lower_bounds(observed_active_counts))
|
|
elif hasattr(modelhub, "account_capacity_limits"):
|
|
account_caps = list(modelhub.account_capacity_limits())
|
|
else:
|
|
default_cap = max(1, int(getattr(modelhub, "active_task_cap", 100) or 100))
|
|
account_caps = [default_cap for _ in clients]
|
|
queue_thresholds = {
|
|
index: max(0, int(account_caps[index]) - reserved_slots)
|
|
for index in range(min(len(clients), len(account_caps)))
|
|
}
|
|
|
|
model_ids = {task.model_id for task in tasks}
|
|
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()}"
|
|
)
|
|
|
|
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)
|
|
}
|
|
block_gpu_task_pairs = {
|
|
(gpu, task_type)
|
|
for gpu, _framework, task_type in block_combinations
|
|
if gpu and task_type
|
|
}
|
|
framework_catalog_combinations: set[tuple[str, str]] = set()
|
|
architecture_model_ids: set[str] = set()
|
|
for task in tasks:
|
|
context = task_contexts.get(str(task.task_id))
|
|
context = context if isinstance(context, dict) else {}
|
|
context_model = str(context.get("modelId") or "").strip()
|
|
context_gpu = str(context.get("targetGpu") or "").strip()
|
|
if (
|
|
(context_model and context_model != task.model_id)
|
|
or (context_gpu and context_gpu.casefold() != task.gpu_type.casefold())
|
|
):
|
|
continue
|
|
framework = str(context.get("framework") or task.framework or "").strip()
|
|
task_type = str(context.get("taskType") or task.task_type or "").strip()
|
|
combination = (
|
|
task.gpu_type.casefold(),
|
|
framework.casefold(),
|
|
task_type.casefold(),
|
|
)
|
|
gpu_task_pair = (task.gpu_type.casefold(), task_type.casefold())
|
|
profile = context.get("modelProfile")
|
|
profile = profile if isinstance(profile, dict) else {}
|
|
exact_relevant = bool(framework and combination in block_combinations)
|
|
exhaustive_relevant = bool(not framework and gpu_task_pair in block_gpu_task_pairs)
|
|
if exhaustive_relevant:
|
|
framework_catalog_combinations.add((task.gpu_type, task_type))
|
|
if (exact_relevant or exhaustive_relevant) and not (
|
|
profile.get("modelType") or profile.get("architectures")
|
|
):
|
|
architecture_model_ids.add(task.model_id)
|
|
framework_catalog, framework_catalog_errors = _load_framework_catalog(
|
|
framework_catalog_combinations,
|
|
modelhub=modelhub,
|
|
read_concurrency=read_concurrency,
|
|
log=log,
|
|
)
|
|
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,
|
|
framework_catalog=framework_catalog,
|
|
)
|
|
log(
|
|
f"[queue-cleanup] architecture_incompatible={len(architecture_decisions)} "
|
|
f"blocks={len(architecture_blocks)} "
|
|
f"rule_state={'ready' if architecture_blocks else 'no_learned_blocks'} "
|
|
f"context_unknown={architecture_skipped['submissionContextUnknown']} "
|
|
f"framework_unknown={architecture_skipped['frameworkContextUnknown']} "
|
|
f"catalog_unknown={architecture_skipped['frameworkCatalogUnknown']} "
|
|
f"partially_blocked={architecture_skipped['partiallyBlockedFrameworkSet']} "
|
|
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 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)
|
|
enriched["cleanupReasons"] = [decision["reason"]]
|
|
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)
|
|
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]] = {}
|
|
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(
|
|
{
|
|
"accountIndex": account_index + 1,
|
|
"error": "active_task_recheck_failed",
|
|
"details": details,
|
|
}
|
|
)
|
|
|
|
by_account: dict[int, list[dict[str, Any]]] = {}
|
|
for decision in decisions:
|
|
account_index = int(decision["accountIndex"]) - 1
|
|
if account_index in refresh_errors:
|
|
continue
|
|
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"}
|
|
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)))
|
|
stop_failed = False
|
|
for account_index in sorted(by_account):
|
|
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"):
|
|
task_ids = sorted(
|
|
task_id
|
|
for task_id, decision in decisions_by_id.items()
|
|
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 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:
|
|
for task in _fetch_status_tasks(
|
|
clients[account_index],
|
|
account_index=account_index,
|
|
status=status,
|
|
):
|
|
phase_tasks[task.task_id] = task
|
|
except Exception as exc:
|
|
stop_errors.append(
|
|
{
|
|
"accountIndex": account_index + 1,
|
|
"taskIds": task_ids,
|
|
"error": f"policy_final_recheck_failed: {type(exc).__name__}: {exc}",
|
|
}
|
|
)
|
|
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")]
|
|
)
|
|
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],
|
|
"recheckedQueuePosition": current_position,
|
|
"recheckedStatus": current_task.status,
|
|
"policyChangeReason": (
|
|
"task_started_running"
|
|
if current_task.status == "running"
|
|
else "queue_position_status_or_policy_changed"
|
|
),
|
|
}
|
|
)
|
|
continue
|
|
decisions_by_id[task_id]["recheckedQueuePosition"] = current_position
|
|
eligible_task_ids.append(task_id)
|
|
task_ids = eligible_task_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()
|
|
|
|
return {
|
|
"dryRun": bool(dry_run),
|
|
"architectureOnly": bool(architecture_only),
|
|
"accounts": len(clients),
|
|
"activeScanned": len(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(oom_decisions),
|
|
"certainOomTasks": oom_decisions,
|
|
"architectureBlockCount": len(architecture_blocks),
|
|
"architectureIncompatibleCount": len(architecture_decisions),
|
|
"architectureIncompatibleTasks": architecture_decisions,
|
|
"architectureModelConfigsComplete": len(model_configs),
|
|
"architectureModelConfigErrors": model_config_errors,
|
|
"architectureFrameworkCatalog": {
|
|
f"{gpu}|{task_type}": sorted(frameworks, key=str.casefold)
|
|
for (gpu, task_type), frameworks in framework_catalog.items()
|
|
},
|
|
"architectureFrameworkCatalogErrors": framework_catalog_errors,
|
|
"architecturePolicySkipped": architecture_skipped,
|
|
"oldOverflowCount": len(old_overflow_decisions),
|
|
"oldOverflowTasks": old_overflow_decisions,
|
|
"oldModelQueueThresholds": [
|
|
queue_thresholds.get(index) for index in range(len(clients))
|
|
],
|
|
"accountCapacityLimits": [
|
|
account_caps[index] if index < len(account_caps) else None
|
|
for index in range(len(clients))
|
|
],
|
|
"recentModelReserveSlots": reserved_slots,
|
|
"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 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)
|
|
parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH))
|
|
parser.add_argument("--modelhub-base-url", default="https://modelhub.org.cn")
|
|
parser.add_argument("--hf-base-url", default="https://modelscope.cn")
|
|
parser.add_argument("--modelhub-token", default=None)
|
|
parser.add_argument("--hf-token", default=None)
|
|
parser.add_argument("--modelscope-token", default=None)
|
|
parser.add_argument("--report-path", default=".modelhub_state/queue_cleanup_latest.json")
|
|
return parser
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
args = build_parser().parse_args(argv)
|
|
ensure_tokens(args)
|
|
tokens = list(getattr(args, "modelhub_tokens", None) or [args.modelhub_token])
|
|
clients = [ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in tokens]
|
|
pool = ModelHubClientPool(clients, capacity_state_path=None)
|
|
discovery = HuggingFaceDiscovery(base_url=args.hf_base_url)
|
|
summary = cleanup_certain_oom_tasks(
|
|
pool,
|
|
discovery,
|
|
dry_run=not args.execute,
|
|
read_concurrency=args.read_concurrency,
|
|
stop_batch_size=args.stop_batch_size,
|
|
)
|
|
report_path = Path(args.report_path)
|
|
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}",
|
|
flush=True,
|
|
)
|
|
return 1 if summary["stopErrors"] else 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|