fix: protect running validation tasks from queue cleanup
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user