fix: protect running validation tasks from queue cleanup

This commit is contained in:
CoolBoy
2026-08-12 01:04:34 +08:00
parent 2065ad6abc
commit d908706f9a
8 changed files with 343 additions and 20 deletions

View File

@@ -261,6 +261,7 @@ def find_old_overflow_tasks(
"recentOverflowTasks": 0,
"modelAgeUnknown": 0,
"accountThresholdUnknown": 0,
"runningOverflowProtected": 0,
}
for account_index, account_tasks in sorted(grouped.items()):
if account_index in incomplete_accounts:
@@ -278,6 +279,9 @@ def find_old_overflow_tasks(
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
@@ -390,6 +394,7 @@ def cleanup_certain_oom_tasks(
(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, age_errors = _load_model_last_modified(
overflow_model_ids,
@@ -410,7 +415,8 @@ def cleanup_certain_oom_tasks(
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"age_unknown={age_skipped['modelAgeUnknown']} "
f"running_protected={age_skipped['runningOverflowProtected']}"
)
decisions_by_key: dict[tuple[int, int], dict[str, Any]] = {}
@@ -444,8 +450,10 @@ def cleanup_certain_oom_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_oom_ids = {
int(decision["taskId"])
@@ -481,14 +489,25 @@ def cleanup_certain_oom_tasks(
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"]))
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}
{
**decision,
"recheckedQueuePosition": current_position,
"recheckedStatus": current_status,
"policyChangeReason": (
"task_started_running"
if current_status == "running"
else "queue_position_or_status_changed"
),
}
)
continue
if current_position is not None:
@@ -501,16 +520,72 @@ def cleanup_certain_oom_tasks(
if stop_failed:
break
decisions_by_id = {int(item["taskId"]): item for item in by_account[account_index]}
phase_ids = [
sorted(
for is_oom_phase in (True, False):
task_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:
if not is_oom_phase and task_ids:
# OOM stops can change actual positions, and a waiting task
# can start running after the account-wide recheck above.
# Re-read this account immediately before its age-only stop.
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"age_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)
if (
account_queue_threshold is None
or current_position is None
or current_position <= account_queue_threshold
or current_task.status != "waiting"
):
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_or_status_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)