fix: protect running validation tasks from queue cleanup
This commit is contained in:
@@ -31,16 +31,22 @@ class FakeQueueClient:
|
||||
*,
|
||||
disappear_on_recheck: bool = False,
|
||||
drop_first_on_recheck: bool = False,
|
||||
promote_on_waiting_read: int | None = None,
|
||||
) -> None:
|
||||
self.records = list(records)
|
||||
self.disappear_on_recheck = disappear_on_recheck
|
||||
self.drop_first_on_recheck = drop_first_on_recheck
|
||||
self.promote_on_waiting_read = promote_on_waiting_read
|
||||
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.promote_on_waiting_read == self.waiting_reads:
|
||||
waiting = [record for record in self.records if record["status"] == "waiting"]
|
||||
if waiting:
|
||||
max(waiting, key=lambda item: int(item["taskId"]))["status"] = "running"
|
||||
if self.disappear_on_recheck and self.waiting_reads >= 2:
|
||||
records: list[dict[str, Any]] = []
|
||||
else:
|
||||
@@ -108,6 +114,7 @@ class QueueCleanupTests(unittest.TestCase):
|
||||
for index in range(1, 192)
|
||||
)
|
||||
tasks.append(OwnedTask(0, 92, "owner/recent", "Iluvatar_bi-100", "waiting"))
|
||||
tasks.append(OwnedTask(0, 93, "owner/old", "Iluvatar_bi-100", "running"))
|
||||
tasks.append(OwnedTask(1, 1192, "owner/recent", "Iluvatar_bi-100", "waiting"))
|
||||
|
||||
selected, skipped = find_old_overflow_tasks(
|
||||
@@ -124,6 +131,93 @@ class QueueCleanupTests(unittest.TestCase):
|
||||
self.assertEqual([91, 1191], [item["taskId"] for item in selected])
|
||||
self.assertEqual([91, 191], [item["queuePosition"] for item in selected])
|
||||
self.assertEqual(2, skipped["recentOverflowTasks"])
|
||||
self.assertEqual(1, skipped["runningOverflowProtected"])
|
||||
|
||||
def test_age_cleanup_never_stops_running_overflow_task(self) -> None:
|
||||
records = [
|
||||
{
|
||||
"taskId": index,
|
||||
"modelId": "owner/old",
|
||||
"gpuType": "Iluvatar_bi-100",
|
||||
"status": "running" if index == 91 else "waiting",
|
||||
}
|
||||
for index in range(1, 92)
|
||||
]
|
||||
client = FakeQueueClient(records)
|
||||
pool = ModelHubClientPool(
|
||||
[client], # type: ignore[list-item]
|
||||
active_task_cap=100,
|
||||
recent_model_reserve_slots=10,
|
||||
)
|
||||
summary = cleanup_certain_oom_tasks(
|
||||
pool,
|
||||
FakeDiscovery(
|
||||
{"owner/old": 1 * GIB},
|
||||
{"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
|
||||
), # type: ignore[arg-type]
|
||||
reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
|
||||
log=lambda _message: None,
|
||||
)
|
||||
|
||||
self.assertEqual(0, summary["oldOverflowCount"])
|
||||
self.assertEqual(1, summary["agePolicySkipped"]["runningOverflowProtected"])
|
||||
self.assertEqual([], client.stopped)
|
||||
|
||||
def test_age_cleanup_recheck_releases_task_that_started_running(self) -> None:
|
||||
records = [
|
||||
{
|
||||
"taskId": index,
|
||||
"modelId": "owner/old",
|
||||
"gpuType": "Iluvatar_bi-100",
|
||||
"status": "waiting",
|
||||
}
|
||||
for index in range(1, 92)
|
||||
]
|
||||
# Read 1 is discovery, read 2 is the account-wide mutation recheck,
|
||||
# and read 3 is the final age-only recheck after the OOM phase.
|
||||
client = FakeQueueClient(records, promote_on_waiting_read=3)
|
||||
pool = ModelHubClientPool(
|
||||
[client], # type: ignore[list-item]
|
||||
active_task_cap=100,
|
||||
recent_model_reserve_slots=10,
|
||||
)
|
||||
summary = cleanup_certain_oom_tasks(
|
||||
pool,
|
||||
FakeDiscovery(
|
||||
{"owner/old": 1 * GIB},
|
||||
{"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
|
||||
), # type: ignore[arg-type]
|
||||
read_concurrency=1,
|
||||
reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
|
||||
log=lambda _message: None,
|
||||
)
|
||||
|
||||
self.assertEqual(1, summary["oldOverflowCount"])
|
||||
self.assertEqual(0, summary["cancelledCount"])
|
||||
self.assertEqual("task_started_running", summary["policyNoLongerAppliesTasks"][0]["policyChangeReason"])
|
||||
self.assertEqual([], client.stopped)
|
||||
|
||||
def test_certain_oom_cleanup_can_still_stop_running_task(self) -> None:
|
||||
client = FakeQueueClient(
|
||||
[
|
||||
{
|
||||
"taskId": 1,
|
||||
"modelId": "owner/large",
|
||||
"gpuType": "Iluvatar_bi-100",
|
||||
"status": "running",
|
||||
}
|
||||
]
|
||||
)
|
||||
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(1, summary["cancelledCount"])
|
||||
self.assertEqual([[1]], client.stopped)
|
||||
|
||||
def test_old_overflow_task_is_not_stopped_if_it_moves_inside_dynamic_threshold(self) -> None:
|
||||
records = [
|
||||
|
||||
Reference in New Issue
Block a user