from __future__ import annotations import argparse from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from datetime import datetime, timedelta from pathlib import Path from typing import Any, Callable, Iterable from candidate_preflight import CandidatePreflightAdvisor, MODEL_LOAD_OVERHEAD from common import 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 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 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, ) 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 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_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, } 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 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, 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.""" 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())}" ) 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, 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 } # OOM tasks are stopped first. Rank the age-policy queue as it will look # after those certain failures are gone, so an old task moving into the # protected first N positions is not over-cancelled. age_rank_tasks = [ task for task in tasks if (task.account_index, task.task_id) not in oom_task_keys ] overflow_model_ids = { task.model_id for account_index in range(len(clients)) if account_index not in listing_errors for task in sorted( (item for item in age_rank_tasks if item.account_index == account_index), key=lambda item: item.task_id, )[queue_thresholds.get(account_index, len(age_rank_tasks)):] } 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']}" ) decisions_by_key: dict[tuple[int, int], dict[str, Any]] = {} for decision in oom_decisions: enriched = dict(decision) enriched["cleanupReasons"] = [decision["reason"]] decisions_by_key[(int(decision["accountIndex"]), int(decision["taskId"]))] = enriched for decision in old_overflow_decisions: key = (int(decision["accountIndex"]), int(decision["taskId"])) existing = decisions_by_key.get(key) if existing is None: enriched = dict(decision) enriched["cleanupReasons"] = [decision["reason"]] decisions_by_key[key] = enriched continue existing["cleanupReasons"].append(decision["reason"]) for field in ("queuePosition", "queueThreshold", "modelLastModified", "recentCutoff"): existing[field] = decision[field] decisions = sorted( decisions_by_key.values(), key=lambda item: (int(item["accountIndex"]), int(item["taskId"])), ) cancelled: list[dict[str, Any]] = [] disappeared: list[dict[str, Any]] = [] policy_no_longer_applies: list[dict[str, Any]] = [] stop_errors: list[dict[str, Any]] = [] if decisions and not dry_run: # Re-read each account immediately before mutation. If any active-state # query fails for that account, fail closed and do not terminate its tasks. refreshed_tasks, refresh_errors = collect_active_tasks(clients, read_concurrency=read_concurrency) active_ids_by_account: dict[int, set[int]] = {} active_positions_by_account: dict[int, dict[int, int]] = {} for task in refreshed_tasks: active_ids_by_account.setdefault(task.account_index, set()).add(task.task_id) for account_index in range(len(clients)): planned_oom_ids = { int(decision["taskId"]) for decision in oom_decisions if int(decision["accountIndex"]) - 1 == account_index } ordered_ids = sorted( task.task_id for task in refreshed_tasks if task.account_index == account_index and task.task_id not in planned_oom_ids ) active_positions_by_account[account_index] = { task_id: position for position, task_id in enumerate(ordered_ids, start=1) } for account_index, details in sorted(refresh_errors.items()): stop_errors.append( { "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"} current_position = active_positions_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 ): policy_no_longer_applies.append( {**decision, "recheckedQueuePosition": current_position} ) continue if current_position is not None: decision["recheckedQueuePosition"] = current_position by_account.setdefault(account_index, []).append(decision) batch_size = max(1, min(100, int(stop_batch_size))) 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]} phase_ids = [ sorted( task_id for task_id, decision in decisions_by_id.items() if ("certain_oom_repository_size_exceeds_gpu_capacity" in decision["cleanupReasons"]) == is_oom_phase ) for is_oom_phase in (True, False) ] for task_ids in phase_ids: for batch in _chunks(task_ids, batch_size): try: clients[account_index].stop_tasks(batch) except Exception as exc: stop_errors.append( { "accountIndex": account_index + 1, "taskIds": batch, "error": f"{type(exc).__name__}: {exc}", } ) stop_failed = True break cancelled.extend(decisions_by_id[task_id] for task_id in batch) log( f"[queue-cleanup] account={account_index + 1:02d} " f"cancelled_batch={len(batch)} cancelled_total={len(cancelled)}" ) if stop_failed: break if cancelled and hasattr(modelhub, "refresh_active_counts"): modelhub.refresh_active_counts() return { "dryRun": bool(dry_run), "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, "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 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"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())