CoolBoy/submmit
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat: dynamically clean incompatible architectures

Browse Source
This commit is contained in:
CoolBoy
2026-08-12 08:19:53 +08:00
parent 615bcad124
commit 7ec875563e
12 changed files with 979 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files

183
tests/test_queue_cleanup.py
View File

@@ -12,10 +12,12 @@ 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,
find_old_overflow_tasks,
)
@@ -72,9 +74,11 @@ class FakeDiscovery:
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]
@@ -85,6 +89,10 @@ class FakeDiscovery:
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:
@@ -103,6 +111,181 @@ class RecordingHttpClient:
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_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_old_models_use_each_accounts_own_capacity_minus_ten_threshold(self) -> None:
now = datetime(2026, 8, 11, tzinfo=timezone.utc)
tasks = [