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())