561 lines
21 KiB
Python
561 lines
21 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
import sys
|
|
from datetime import datetime, timezone
|
|
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 architecture_compatibility import architecture_compatibility_key # noqa: E402
|
|
from modelhub_client import ModelHubClient, ModelHubClientPool # noqa: E402
|
|
from queue_cleanup import ( # noqa: E402
|
|
OwnedTask,
|
|
cleanup_certain_oom_tasks,
|
|
find_architecture_incompatible_tasks,
|
|
find_certain_oom_tasks,
|
|
)
|
|
|
|
|
|
GIB = 1024**3
|
|
|
|
|
|
class FakeQueueClient:
|
|
def __init__(
|
|
self,
|
|
records: list[dict[str, Any]],
|
|
*,
|
|
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:
|
|
records = [record for record in self.records if record["status"] == "waiting"]
|
|
if self.drop_first_on_recheck and self.waiting_reads >= 2:
|
|
records = sorted(records, key=lambda item: int(item["taskId"]))[1:]
|
|
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 list_framework_stats(self, task_type: str, target_gpu: str) -> list[dict[str, Any]]:
|
|
del task_type, target_gpu
|
|
return [
|
|
{"framework": "vllm", "modelCount": 100, "successCount": 60},
|
|
{"framework": "transformers", "modelCount": 100, "successCount": 50},
|
|
]
|
|
|
|
def count_active_tasks(self, **_kwargs: Any) -> int:
|
|
return len(self.records)
|
|
|
|
|
|
class FakeDiscovery:
|
|
def __init__(
|
|
self,
|
|
sizes: dict[str, int | None],
|
|
last_modified: dict[str, datetime | None] | None = None,
|
|
configs: dict[str, dict[str, Any]] | None = None,
|
|
) -> None:
|
|
self.sizes = sizes
|
|
self.last_modified = last_modified or {}
|
|
self.configs = configs or {}
|
|
|
|
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}]
|
|
|
|
def get_model_last_modified(self, repo_id: str) -> datetime | None:
|
|
return self.last_modified.get(repo_id)
|
|
|
|
def get_model_config(self, repo_id: str) -> tuple[dict[str, Any], str | None]:
|
|
config = self.configs.get(repo_id)
|
|
return (dict(config), None) if config is not None else ({}, "config_not_found")
|
|
|
|
|
|
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):
|
|
@staticmethod
|
|
def architecture_block(
|
|
*,
|
|
gpu: str = "Iluvatar_bi-100",
|
|
framework: str = "vllm",
|
|
task_type: str = "text-generation",
|
|
signature: str = "architectures:qwen2forcausallm",
|
|
) -> tuple[str, dict[str, Any]]:
|
|
key = architecture_compatibility_key(gpu, framework, task_type, signature)
|
|
assert key is not None
|
|
return key, {
|
|
"targetGpu": gpu,
|
|
"framework": framework,
|
|
"taskType": task_type,
|
|
"matchType": "architectures",
|
|
"architectureSignature": signature,
|
|
"evidenceCount": 1,
|
|
"expiresAt": "2026-09-11T00:00:00+00:00",
|
|
}
|
|
|
|
def test_architecture_cleanup_matches_exact_context_and_protects_running(self) -> None:
|
|
key, block = self.architecture_block()
|
|
tasks = [
|
|
OwnedTask(0, 1, "owner/waiting", "Iluvatar_bi-100", "waiting"),
|
|
OwnedTask(0, 2, "owner/running", "Iluvatar_bi-100", "running"),
|
|
OwnedTask(0, 3, "owner/other-framework", "Iluvatar_bi-100", "waiting"),
|
|
]
|
|
contexts = {
|
|
"1": {
|
|
"modelId": "owner/waiting",
|
|
"targetGpu": "Iluvatar_bi-100",
|
|
"framework": "vllm",
|
|
"taskType": "text-generation",
|
|
"modelProfile": {"architectures": ["Qwen2ForCausalLM"]},
|
|
},
|
|
"2": {
|
|
"modelId": "owner/running",
|
|
"targetGpu": "Iluvatar_bi-100",
|
|
"framework": "vllm",
|
|
"taskType": "text-generation",
|
|
"modelProfile": {"architectures": ["Qwen2ForCausalLM"]},
|
|
},
|
|
"3": {
|
|
"modelId": "owner/other-framework",
|
|
"targetGpu": "Iluvatar_bi-100",
|
|
"framework": "mindie",
|
|
"taskType": "text-generation",
|
|
"modelProfile": {"architectures": ["Qwen2ForCausalLM"]},
|
|
},
|
|
}
|
|
|
|
selected, skipped = find_architecture_incompatible_tasks(
|
|
tasks,
|
|
architecture_blocks={key: block},
|
|
task_contexts=contexts,
|
|
model_configs={},
|
|
)
|
|
|
|
self.assertEqual([1], [item["taskId"] for item in selected])
|
|
self.assertEqual(1, skipped["runningMatchedProtected"])
|
|
self.assertEqual(1, skipped["noMatchingBlock"])
|
|
|
|
def test_architecture_cleanup_uses_current_queue_task_type_when_every_framework_is_blocked(self) -> None:
|
|
vllm_key, vllm_block = self.architecture_block(framework="vllm")
|
|
transformers_key, transformers_block = self.architecture_block(framework="transformers")
|
|
tasks = [
|
|
OwnedTask(
|
|
0,
|
|
1,
|
|
"owner/model",
|
|
"Iluvatar_bi-100",
|
|
"waiting",
|
|
task_type="text-generation",
|
|
)
|
|
]
|
|
|
|
selected, skipped = find_architecture_incompatible_tasks(
|
|
tasks,
|
|
architecture_blocks={
|
|
vllm_key: vllm_block,
|
|
transformers_key: transformers_block,
|
|
},
|
|
task_contexts={},
|
|
model_configs={
|
|
"owner/model": {"architectures": ["Qwen2ForCausalLM"]},
|
|
},
|
|
framework_catalog={
|
|
("iluvatar_bi-100", "text-generation"): {"vllm", "transformers"},
|
|
},
|
|
)
|
|
|
|
self.assertEqual([1], [item["taskId"] for item in selected])
|
|
self.assertEqual("all_live_frameworks", selected[0]["architectureMatchScope"])
|
|
self.assertEqual(["transformers", "vllm"], selected[0]["evaluatedFrameworks"])
|
|
self.assertEqual(1, skipped["frameworkContextUnknown"])
|
|
|
|
def test_architecture_cleanup_does_not_guess_when_only_some_frameworks_are_blocked(self) -> None:
|
|
key, block = self.architecture_block(framework="vllm")
|
|
tasks = [
|
|
OwnedTask(
|
|
0,
|
|
1,
|
|
"owner/model",
|
|
"Iluvatar_bi-100",
|
|
"waiting",
|
|
task_type="text-generation",
|
|
)
|
|
]
|
|
|
|
selected, skipped = find_architecture_incompatible_tasks(
|
|
tasks,
|
|
architecture_blocks={key: block},
|
|
task_contexts={},
|
|
model_configs={
|
|
"owner/model": {"architectures": ["Qwen2ForCausalLM"]},
|
|
},
|
|
framework_catalog={
|
|
("iluvatar_bi-100", "text-generation"): {"vllm", "transformers"},
|
|
},
|
|
)
|
|
|
|
self.assertEqual([], selected)
|
|
self.assertEqual(1, skipped["partiallyBlockedFrameworkSet"])
|
|
|
|
def test_queue_cleanup_enriches_active_api_record_and_stops_only_all_framework_block(self) -> None:
|
|
vllm_key, vllm_block = self.architecture_block(framework="vllm")
|
|
transformers_key, transformers_block = self.architecture_block(framework="transformers")
|
|
client = FakeQueueClient(
|
|
[
|
|
{
|
|
"taskId": 1,
|
|
"modelId": "owner/model",
|
|
"gpuType": "Iluvatar_bi-100",
|
|
"status": "waiting",
|
|
"modelTaskLevelId": 23,
|
|
"modelTaskLevel": "文本生成",
|
|
}
|
|
]
|
|
)
|
|
pool = ModelHubClientPool([client], active_task_cap=100) # type: ignore[list-item]
|
|
|
|
summary = cleanup_certain_oom_tasks(
|
|
pool,
|
|
FakeDiscovery(
|
|
{"owner/model": 1 * GIB},
|
|
configs={"owner/model": {"architectures": ["Qwen2ForCausalLM"]}},
|
|
), # type: ignore[arg-type]
|
|
architecture_compatibility_blocks={
|
|
vllm_key: vllm_block,
|
|
transformers_key: transformers_block,
|
|
},
|
|
task_compatibility_contexts={},
|
|
log=lambda _message: None,
|
|
)
|
|
|
|
self.assertEqual(1, summary["architectureIncompatibleCount"])
|
|
self.assertEqual(1, summary["cancelledCount"])
|
|
self.assertEqual(
|
|
["transformers", "vllm"],
|
|
summary["architectureFrameworkCatalog"]["iluvatar_bi-100|text-generation"],
|
|
)
|
|
self.assertEqual([[1]], client.stopped)
|
|
|
|
def test_queue_cleanup_fetches_config_and_stops_known_incompatible_waiting_task(self) -> None:
|
|
key, block = self.architecture_block()
|
|
client = FakeQueueClient(
|
|
[
|
|
{
|
|
"taskId": 1,
|
|
"modelId": "owner/model",
|
|
"gpuType": "Iluvatar_bi-100",
|
|
"status": "waiting",
|
|
}
|
|
]
|
|
)
|
|
pool = ModelHubClientPool([client], active_task_cap=100) # type: ignore[list-item]
|
|
summary = cleanup_certain_oom_tasks(
|
|
pool,
|
|
FakeDiscovery(
|
|
{"owner/model": 1 * GIB},
|
|
configs={
|
|
"owner/model": {
|
|
"model_type": "qwen2",
|
|
"architectures": ["Qwen2ForCausalLM"],
|
|
}
|
|
},
|
|
), # type: ignore[arg-type]
|
|
architecture_compatibility_blocks={key: block},
|
|
task_compatibility_contexts={
|
|
"1": {
|
|
"modelId": "owner/model",
|
|
"targetGpu": "Iluvatar_bi-100",
|
|
"framework": "vllm",
|
|
"taskType": "text-generation",
|
|
"modelProfile": {},
|
|
}
|
|
},
|
|
log=lambda _message: None,
|
|
)
|
|
|
|
self.assertEqual(1, summary["architectureIncompatibleCount"])
|
|
self.assertEqual(1, summary["cancelledCount"])
|
|
self.assertEqual([[1]], client.stopped)
|
|
|
|
def test_dynamic_architecture_only_cleanup_skips_expensive_size_and_age_scans(self) -> None:
|
|
key, block = self.architecture_block()
|
|
client = FakeQueueClient(
|
|
[
|
|
{
|
|
"taskId": 1,
|
|
"modelId": "owner/model",
|
|
"gpuType": "Iluvatar_bi-100",
|
|
"status": "waiting",
|
|
}
|
|
]
|
|
)
|
|
pool = ModelHubClientPool([client], active_task_cap=100) # type: ignore[list-item]
|
|
summary = cleanup_certain_oom_tasks(
|
|
pool,
|
|
FakeDiscovery({}), # type: ignore[arg-type]
|
|
architecture_compatibility_blocks={key: block},
|
|
task_compatibility_contexts={
|
|
"1": {
|
|
"modelId": "owner/model",
|
|
"targetGpu": "Iluvatar_bi-100",
|
|
"framework": "vllm",
|
|
"taskType": "text-generation",
|
|
"modelProfile": {"architectures": ["Qwen2ForCausalLM"]},
|
|
}
|
|
},
|
|
architecture_only=True,
|
|
log=lambda _message: None,
|
|
)
|
|
|
|
self.assertTrue(summary["architectureOnly"])
|
|
self.assertEqual(0, summary["repositorySizesComplete"])
|
|
self.assertEqual(0, summary["modelAgeMetadataComplete"])
|
|
self.assertEqual(1, summary["architectureIncompatibleCount"])
|
|
self.assertEqual([[1]], client.stopped)
|
|
|
|
def test_architecture_cleanup_recheck_protects_task_that_started_running(self) -> None:
|
|
key, block = self.architecture_block()
|
|
client = FakeQueueClient(
|
|
[
|
|
{
|
|
"taskId": 1,
|
|
"modelId": "owner/model",
|
|
"gpuType": "Iluvatar_bi-100",
|
|
"status": "waiting",
|
|
}
|
|
],
|
|
promote_on_waiting_read=2,
|
|
)
|
|
pool = ModelHubClientPool([client], active_task_cap=100) # type: ignore[list-item]
|
|
summary = cleanup_certain_oom_tasks(
|
|
pool,
|
|
FakeDiscovery({"owner/model": 1 * GIB}), # type: ignore[arg-type]
|
|
architecture_compatibility_blocks={key: block},
|
|
task_compatibility_contexts={
|
|
"1": {
|
|
"modelId": "owner/model",
|
|
"targetGpu": "Iluvatar_bi-100",
|
|
"framework": "vllm",
|
|
"taskType": "text-generation",
|
|
"modelProfile": {"architectures": ["Qwen2ForCausalLM"]},
|
|
}
|
|
},
|
|
read_concurrency=1,
|
|
log=lambda _message: None,
|
|
)
|
|
|
|
self.assertEqual(1, summary["architectureIncompatibleCount"])
|
|
self.assertEqual(0, summary["cancelledCount"])
|
|
self.assertEqual("task_started_running", summary["policyNoLongerAppliesTasks"][0]["policyChangeReason"])
|
|
self.assertEqual([], client.stopped)
|
|
|
|
def test_initial_and_later_cleanup_never_stop_old_waiting_or_running_tasks(self) -> None:
|
|
records = [
|
|
{
|
|
"taskId": index,
|
|
"modelId": "owner/old",
|
|
"gpuType": "Iluvatar_bi-100",
|
|
"status": "running" if index == 100 else "waiting",
|
|
}
|
|
for index in range(1, 101)
|
|
]
|
|
client = FakeQueueClient(records)
|
|
pool = ModelHubClientPool(
|
|
[client], # type: ignore[list-item]
|
|
active_task_cap=100,
|
|
recent_model_reserve_slots=5,
|
|
)
|
|
summaries = [
|
|
cleanup_certain_oom_tasks(
|
|
pool,
|
|
FakeDiscovery({"owner/old": 1 * GIB}), # type: ignore[arg-type]
|
|
log=lambda _message: None,
|
|
)
|
|
for _ in range(2)
|
|
]
|
|
|
|
self.assertTrue(all(item["ageCleanupMode"] == "admission_only" for item in summaries))
|
|
self.assertTrue(all(item["oldOverflowCount"] == 0 for item in summaries))
|
|
self.assertTrue(all(item["cancelledCount"] == 0 for item in summaries))
|
|
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_capacity_decline_pauses_submissions_and_never_deletes_existing_tasks(self) -> None:
|
|
records = [
|
|
{
|
|
"taskId": index,
|
|
"modelId": "owner/old",
|
|
"gpuType": "Iluvatar_bi-100",
|
|
"status": "waiting",
|
|
}
|
|
for index in range(1, 151)
|
|
]
|
|
client = FakeQueueClient(records)
|
|
pool = ModelHubClientPool(
|
|
[client], # type: ignore[list-item]
|
|
active_task_cap=100,
|
|
recent_model_reserve_slots=5,
|
|
capacity_state_path=None,
|
|
)
|
|
summary = cleanup_certain_oom_tasks(
|
|
pool,
|
|
FakeDiscovery({"owner/old": 1 * GIB}), # type: ignore[arg-type]
|
|
log=lambda _message: None,
|
|
)
|
|
|
|
self.assertEqual([150], summary["accountCapacityLimits"])
|
|
self.assertEqual([145], summary["oldModelQueueThresholds"])
|
|
self.assertEqual(0, summary["oldOverflowCount"])
|
|
self.assertEqual(0, pool.available_submit_slots())
|
|
self.assertEqual([], client.stopped)
|
|
|
|
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()
|