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

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