feat: clean deterministic OOM tasks on startup

This commit is contained in:
CoolBoy
2026-08-11 00:52:16 +08:00
parent 3d15f60284
commit 9645973468
7 changed files with 641 additions and 1 deletions

View File

@@ -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

View File

@@ -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

View File

@@ -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:

View File

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

View File

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

View File

@@ -1 +1 @@
AGENT_VERSION = "2026.08.10.3"
AGENT_VERSION = "2026.08.11.1"

160
tests/test_queue_cleanup.py Normal file
View File

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