From 964597346838b0f89bdfc70a7195bb174e5bd062 Mon Sep 17 00:00:00 2001 From: CoolBoy Date: Tue, 11 Aug 2026 00:52:16 +0800 Subject: [PATCH] feat: clean deterministic OOM tasks on startup --- README.md | 19 ++ modelhub_submmit_api/README.md | 7 + modelhub_submmit_api/modelhub_client.py | 28 ++ modelhub_submmit_api/poll_runner.py | 62 ++++ modelhub_submmit_api/queue_cleanup.py | 364 ++++++++++++++++++++++++ modelhub_submmit_api/version.py | 2 +- tests/test_queue_cleanup.py | 160 +++++++++++ 7 files changed, 641 insertions(+), 1 deletion(-) create mode 100644 modelhub_submmit_api/queue_cleanup.py create mode 100644 tests/test_queue_cleanup.py diff --git a/README.md b/README.md index 35beb60d..31d297a8 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,9 @@ Optional tuning: - `MODELSCOPE_PAGE_INTERVAL_SECONDS` default `0.25` - `MODELSCOPE_PAGE_CACHE_TTL_SECONDS` default `900` - `MODELHUB_AGENT_VERIFY_CACHE_TTL_SECONDS` default `900` +- `MODELHUB_QUEUE_CLEANUP_INTERVAL_CYCLES` default `120`; cleanup also runs once at startup +- `MODELHUB_QUEUE_CLEANUP_READ_CONCURRENCY` default `6` +- `MODELHUB_QUEUE_CLEANUP_REPORT_PATH` default `.modelhub_state/queue_cleanup_latest.json` ## Adaptive GPU Strategy @@ -119,6 +122,19 @@ capacities with `MODELHUB_GPU_MEMORY_GIB_JSON`, for example `{"New_gpu": 64}`. +At poller startup, the same deterministic memory gate is applied to existing +`waiting` and `running` tasks across every configured account. A task is stopped +through `PUT /api/async/task/stop-create-contest-task` only when its own current +recursive repository size, multiplied by ModelHub's observed `1.20` overhead, +exceeds the known capacity of its selected GPU. The task ID is fetched and +stopped with the token belonging to that account, and its active state is +rechecked immediately before the mutation. Missing file sizes, unknown GPU +capacities, listing failures, and tasks that have already changed state all fail +closed and are never stopped. This does not match against another task from the +same model or infer failure from historical similarity. The cleanup repeats +every 120 poll cycles by default and writes its full evidence report to +`.modelhub_state/queue_cleanup_latest.json`. + The verified capacities, safe repository-size boundaries, evidence hierarchy, and source links are recorded in `docs/gpu-memory-capacity-2026-08-10.md`. @@ -207,6 +223,9 @@ GPU, recursive repository-size checks, deterministic failure-aware preflight, and rate-limited lazy Qwen review for unresolved semantic cases. Version `2026.08.10.3` selects `qwen3.7-flash` by default and recognizes the repository root `.env` key named `dashscope` without logging its value. +Version `2026.08.11.1` adds account-owned cancellation of queued tasks that are +deterministically over the selected GPU's ModelHub memory boundary, with a +second active-state check and fail-closed handling for incomplete evidence. ## Deploy diff --git a/modelhub_submmit_api/README.md b/modelhub_submmit_api/README.md index 8feae698..775397de 100644 --- a/modelhub_submmit_api/README.md +++ b/modelhub_submmit_api/README.md @@ -16,6 +16,7 @@ It currently supports: - `main.py`: core discovery, scoring, dedup, and submission - `daily_runner.py`: daily wave orchestration - `poll_runner.py`: long-running queue refiller +- `queue_cleanup.py`: fail-closed cleanup for active tasks that are certain to exceed GPU memory - `runner_common.py`: shared token / key file loading - `hf_discovery.py`: ModelScope model discovery and inspection (keeps the legacy module name) - `modelhub_client.py`: ModelHub API client and token-pool routing @@ -115,6 +116,10 @@ bash run_poll.sh --dry-run raises that account's persisted known limit; a capacity rejection enters cooldown. - Each `[scan]` log records the discovery stage and candidate yield. The final `[daily] wave_done` log includes `skip_reasons`, making empty candidate pools distinguishable from API failures. +- At startup and every 120 poll cycles, active tasks are checked with the same recursive-size + memory rule as new submissions. Only tasks whose own repository size times `1.20` exceeds + their selected GPU capacity are stopped, after a fresh account-scoped active-state check. + Incomplete size/capacity evidence is never used for cancellation. ## Important Flags @@ -148,6 +153,7 @@ Common flags: - `--submit-concurrency`: concurrent task submission calls used by each cycle (0 = auto) - `--post-cycle-cooldown-seconds`: pause after a successful cycle before next cycle (default 2) - `--max-cycles`: optional hard stop for testing or batch windows +- `--disable-queue-cleanup`: disable automatic deterministic OOM cleanup Failure-informed preflight is enabled by default. It rejects deterministic missing-file and predicted-OOM cases, clamps unsafe context-length arguments, @@ -198,6 +204,7 @@ Persistent local scheduler state is written under `.modelhub_state/`: - `market_intelligence.json`: cached public queue, throughput, health, and framework statistics - `account_capacity.json`: learned per-account active-task limits - `submission_exclusions.jsonl`: non-retryable model/GPU uniqueness rejections +- `queue_cleanup_latest.json`: latest active-task sizing evidence and cancellation result ## Verification diff --git a/modelhub_submmit_api/modelhub_client.py b/modelhub_submmit_api/modelhub_client.py index 77862723..fc84540f 100644 --- a/modelhub_submmit_api/modelhub_client.py +++ b/modelhub_submmit_api/modelhub_client.py @@ -176,6 +176,28 @@ class ModelHubClient: def add_task(self, payload: dict[str, Any]) -> dict[str, Any]: return self._request("POST", "/api/adapt/task/add", data=payload) + def stop_tasks(self, task_ids: list[int | str]) -> dict[str, Any]: + """Stop active validation tasks owned by this authenticated account.""" + normalized: list[int] = [] + seen: set[int] = set() + for task_id in task_ids: + try: + numeric_id = int(str(task_id).strip()) + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid ModelHub task ID: {task_id!r}") from exc + if numeric_id <= 0: + raise ValueError(f"Invalid ModelHub task ID: {task_id!r}") + if numeric_id not in seen: + seen.add(numeric_id) + normalized.append(numeric_id) + if not normalized: + raise ValueError("At least one ModelHub task ID is required") + return self._request( + "PUT", + "/api/async/task/stop-create-contest-task", + data={"taskIds": normalized}, + ) + def find_recent_task_id(self, model_id: str, gpu_type: str, submitted_after: datetime) -> str | None: recent_tasks = self.list_tasks( page_size=20, @@ -560,6 +582,12 @@ class ModelHubClientPool: with self._state_lock: return self._counts_snapshot_locked() + def refresh_active_counts(self) -> list[int]: + """Force a remote refresh after out-of-band queue mutations.""" + self._refresh_active_counts(force=True) + with self._state_lock: + return self._counts_snapshot_locked() + def available_submit_slots(self) -> int: self._refresh_active_counts() with self._state_lock: diff --git a/modelhub_submmit_api/poll_runner.py b/modelhub_submmit_api/poll_runner.py index b1485056..b131502e 100644 --- a/modelhub_submmit_api/poll_runner.py +++ b/modelhub_submmit_api/poll_runner.py @@ -23,6 +23,7 @@ from market_intelligence import ( ) from modelhub_client import DEFAULT_CAPACITY_STATE_PATH, ModelHubClient, ModelHubClientPool from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker +from queue_cleanup import cleanup_certain_oom_tasks from runner_common import DEFAULT_KEY_PATH, ensure_tokens from submission_claims import DEFAULT_CLAIMS_PATH from template_selector import TemplateSelector @@ -167,6 +168,24 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--post-cycle-cooldown-seconds", type=int, default=2, help="Short sleep after a successful cycle") parser.add_argument("--max-cycles", type=int, default=0, help="Optional hard stop after N cycles; 0 means run until quota is reached") parser.add_argument("--print-stats", action="store_true", help="Load outcomes, sync, print stats report, and exit") + parser.add_argument("--disable-queue-cleanup", action="store_true", help=argparse.SUPPRESS) + parser.add_argument( + "--queue-cleanup-interval-cycles", + type=int, + default=int(os.getenv("MODELHUB_QUEUE_CLEANUP_INTERVAL_CYCLES", "120")), + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--queue-cleanup-read-concurrency", + type=int, + default=int(os.getenv("MODELHUB_QUEUE_CLEANUP_READ_CONCURRENCY", "6")), + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--queue-cleanup-report-path", + default=os.getenv("MODELHUB_QUEUE_CLEANUP_REPORT_PATH", ".modelhub_state/queue_cleanup_latest.json"), + help=argparse.SUPPRESS, + ) return parser @@ -220,6 +239,7 @@ def run_poll_loop( ) cycle_summaries: list[dict[str, Any]] = [] + queue_cleanup_runs: list[dict[str, Any]] = [] submitted_total = 0 cycles = 0 stopped_reason = "max_cycles_reached" @@ -232,6 +252,47 @@ def run_poll_loop( cycles += 1 if hasattr(modelhub_client, "configure_capacity_probe"): modelhub_client.configure_capacity_probe(cycles) + + cleanup_interval = max(0, int(getattr(base_args, "queue_cleanup_interval_cycles", 120) or 0)) + should_cleanup_queue = ( + not bool(getattr(base_args, "disable_queue_cleanup", False)) + and isinstance(modelhub_client, ModelHubClientPool) + and (cycles == 1 or (cleanup_interval > 0 and cycles % cleanup_interval == 0)) + ) + if should_cleanup_queue: + try: + cleanup_feedback = outcome_tracker.get_stats_report() + cleanup_gpu_memory = cleanup_feedback.get("observedGpuMemoryGiB") or {} + cleanup_summary = cleanup_certain_oom_tasks( + modelhub_client, + hf_discovery, + dry_run=bool(base_args.dry_run), + read_concurrency=max(1, int(getattr(base_args, "queue_cleanup_read_concurrency", 6) or 6)), + gpu_memory_gib=cleanup_gpu_memory if isinstance(cleanup_gpu_memory, dict) else None, + log=log, + ) + write_json( + Path( + getattr( + base_args, + "queue_cleanup_report_path", + ".modelhub_state/queue_cleanup_latest.json", + ) + ), + cleanup_summary, + ) + queue_cleanup_runs.append( + { + "cycle": cycles, + "activeScanned": cleanup_summary["activeScanned"], + "certainOomCount": cleanup_summary["certainOomCount"], + "cancelledCount": cleanup_summary["cancelledCount"], + "stopErrorCount": len(cleanup_summary["stopErrors"]), + } + ) + except Exception as exc: + log(f"[queue-cleanup] error={type(exc).__name__}: {exc} continue_polling=true") + active_counts = modelhub_client.active_task_counts() if hasattr(modelhub_client, "active_task_counts") else [] capacity_limits = modelhub_client.account_capacity_limits() if hasattr(modelhub_client, "account_capacity_limits") else [] capacity_probe = modelhub_client.capacity_probe_enabled() if hasattr(modelhub_client, "capacity_probe_enabled") else False @@ -335,6 +396,7 @@ def run_poll_loop( "stoppedReason": stopped_reason, "pollRunDir": str(poll_run_dir), "cycleSummaries": cycle_summaries, + "queueCleanupRuns": queue_cleanup_runs, "outcomeStats": stats_report, } write_json(poll_run_dir / "summary.json", summary) diff --git a/modelhub_submmit_api/queue_cleanup.py b/modelhub_submmit_api/queue_cleanup.py new file mode 100644 index 00000000..e4640b21 --- /dev/null +++ b/modelhub_submmit_api/queue_cleanup.py @@ -0,0 +1,364 @@ +from __future__ import annotations + +import argparse +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Iterable + +from candidate_preflight import CandidatePreflightAdvisor, MODEL_LOAD_OVERHEAD +from common import write_json +from hf_discovery import HuggingFaceDiscovery, inspect_repo_tree +from modelhub_client import ModelHubClient, ModelHubClientPool +from runner_common import DEFAULT_KEY_PATH, ensure_tokens + + +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 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 _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, + log: Callable[[str], None] = print, +) -> dict[str, Any]: + """Stop only active tasks that cannot fit the selected GPU by known capacity.""" + clients = list(modelhub.clients) + 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())}" + ) + + 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, + ) + decisions, skipped = find_certain_oom_tasks( + tasks, + repository_sizes=repository_sizes, + gpu_memory_gib=gpu_memory_gib, + ) + log( + f"[queue-cleanup] certain_oom={len(decisions)} " + f"fits={skipped['fitsKnownCapacity']} size_unknown={skipped['repositorySizeUnknown']} " + f"gpu_unknown={skipped['gpuCapacityUnknown']} dry_run={str(bool(dry_run)).lower()}" + ) + + cancelled: list[dict[str, Any]] = [] + disappeared: 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]] = {} + for task in refreshed_tasks: + active_ids_by_account.setdefault(task.account_index, set()).add(task.task_id) + + 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 + 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]} + task_ids = sorted(decisions_by_id) + 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 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, + "listingErrors": {str(index + 1): values for index, values in listing_errors.items()}, + "certainOomCount": len(decisions), + "certainOomTasks": decisions, + "skipped": skipped, + "cancelledCount": len(cancelled), + "cancelledTasks": cancelled, + "noLongerActiveCount": len(disappeared), + "stopErrors": stop_errors, + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Safely stop active ModelHub tasks that are certain to run out of GPU memory.") + 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"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()) diff --git a/modelhub_submmit_api/version.py b/modelhub_submmit_api/version.py index e7ecaa49..d82e8566 100644 --- a/modelhub_submmit_api/version.py +++ b/modelhub_submmit_api/version.py @@ -1 +1 @@ -AGENT_VERSION = "2026.08.10.3" +AGENT_VERSION = "2026.08.11.1" diff --git a/tests/test_queue_cleanup.py b/tests/test_queue_cleanup.py new file mode 100644 index 00000000..39db92ca --- /dev/null +++ b/tests/test_queue_cleanup.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import unittest +import sys +from pathlib import Path +from typing import Any + + +PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api" +if str(PACKAGE_DIR) in sys.path: + sys.path.remove(str(PACKAGE_DIR)) +sys.path.insert(0, str(PACKAGE_DIR)) + +from modelhub_client import ModelHubClient, ModelHubClientPool # noqa: E402 +from queue_cleanup import OwnedTask, cleanup_certain_oom_tasks, find_certain_oom_tasks # noqa: E402 + + +GIB = 1024**3 + + +class FakeQueueClient: + def __init__(self, records: list[dict[str, Any]], *, disappear_on_recheck: bool = False) -> None: + self.records = list(records) + self.disappear_on_recheck = disappear_on_recheck + self.waiting_reads = 0 + self.stopped: list[list[int]] = [] + + def list_tasks_page(self, *, status: str, **_kwargs: Any) -> dict[str, Any]: + if status == "waiting": + self.waiting_reads += 1 + if self.disappear_on_recheck and self.waiting_reads >= 2: + records: list[dict[str, Any]] = [] + else: + records = [record for record in self.records if record["status"] == "waiting"] + else: + records = [record for record in self.records if record["status"] == status] + return {"code": 0, "data": {"records": records, "pages": 1}} + + def stop_tasks(self, task_ids: list[int]) -> dict[str, Any]: + self.stopped.append(list(task_ids)) + ids = set(task_ids) + self.records = [record for record in self.records if int(record["taskId"]) not in ids] + return {"code": 0, "data": None} + + def count_active_tasks(self, **_kwargs: Any) -> int: + return len(self.records) + + +class FakeDiscovery: + def __init__(self, sizes: dict[str, int | None]) -> None: + self.sizes = sizes + + def list_repo_tree(self, repo_id: str) -> list[dict[str, Any]]: + size = self.sizes[repo_id] + if size is None: + return [{"Path": "model.safetensors"}] + return [{"Path": "model.safetensors", "Size": size}] + + +class RecordingHttpClient: + def __init__(self) -> None: + self.calls: list[tuple[str, str, dict[str, Any] | None, dict[str, Any] | None]] = [] + + def request_json( + self, + method: str, + path: str, + *, + query: dict[str, Any] | None = None, + data: dict[str, Any] | None = None, + ) -> dict[str, Any]: + self.calls.append((method, path, query, data)) + return {"code": 0, "data": None} + + +class QueueCleanupTests(unittest.TestCase): + def test_stop_tasks_uses_documented_put_endpoint_and_integer_ids(self) -> None: + http = RecordingHttpClient() + client = ModelHubClient(http_client=http) # type: ignore[arg-type] + + client.stop_tasks(["12", 12, 13]) + + self.assertEqual( + [ + ( + "PUT", + "/api/async/task/stop-create-contest-task", + None, + {"taskIds": [12, 13]}, + ) + ], + http.calls, + ) + + def test_only_exact_size_capacity_failures_are_selected(self) -> None: + tasks = [ + OwnedTask(0, 1, "owner/too-large", "Iluvatar_bi-100", "waiting"), + OwnedTask(0, 2, "owner/fits", "Iluvatar_bi-100", "waiting"), + OwnedTask(0, 3, "owner/unknown-size", "Iluvatar_bi-100", "waiting"), + OwnedTask(0, 4, "owner/unknown-gpu", "new-gpu", "waiting"), + ] + + selected, skipped = find_certain_oom_tasks( + tasks, + repository_sizes={ + "owner/too-large": 30 * GIB, + "owner/fits": 20 * GIB, + "owner/unknown-gpu": 30 * GIB, + }, + ) + + self.assertEqual([1], [item["taskId"] for item in selected]) + self.assertEqual(1, skipped["repositorySizeUnknown"]) + self.assertEqual(1, skipped["gpuCapacityUnknown"]) + self.assertEqual(1, skipped["fitsKnownCapacity"]) + + def test_cleanup_stops_only_certain_oom_tasks_on_the_owning_account(self) -> None: + first = FakeQueueClient( + [ + {"taskId": 1, "modelId": "owner/large", "gpuType": "Iluvatar_bi-100", "status": "waiting"}, + {"taskId": 2, "modelId": "owner/small", "gpuType": "Iluvatar_bi-100", "status": "waiting"}, + ] + ) + second = FakeQueueClient( + [{"taskId": 3, "modelId": "owner/large", "gpuType": "MetaX_c-500", "status": "waiting"}] + ) + pool = ModelHubClientPool([first, second], active_task_cap=100) # type: ignore[list-item] + + summary = cleanup_certain_oom_tasks( + pool, + FakeDiscovery({"owner/large": 40 * GIB, "owner/small": 20 * GIB}), # type: ignore[arg-type] + log=lambda _message: None, + ) + + self.assertEqual(1, summary["certainOomCount"]) + self.assertEqual(1, summary["cancelledCount"]) + self.assertEqual([[1]], first.stopped) + self.assertEqual([], second.stopped) + + def test_task_that_disappears_during_scan_is_not_stopped(self) -> None: + client = FakeQueueClient( + [{"taskId": 1, "modelId": "owner/large", "gpuType": "Iluvatar_bi-100", "status": "waiting"}], + disappear_on_recheck=True, + ) + pool = ModelHubClientPool([client], active_task_cap=100) # type: ignore[list-item] + + summary = cleanup_certain_oom_tasks( + pool, + FakeDiscovery({"owner/large": 40 * GIB}), # type: ignore[arg-type] + log=lambda _message: None, + ) + + self.assertEqual(1, summary["certainOomCount"]) + self.assertEqual(0, summary["cancelledCount"]) + self.assertEqual(1, summary["noLongerActiveCount"]) + self.assertEqual([], client.stopped) + + +if __name__ == "__main__": + unittest.main()