feat: reserve queue capacity for recent models
This commit is contained in:
@@ -17,8 +17,13 @@ if str(PACKAGE_DIR) in sys.path:
|
||||
sys.path.insert(0, str(PACKAGE_DIR))
|
||||
|
||||
from common import read_json, read_jsonl # noqa: E402
|
||||
from main import build_parser, make_run_dir, run_submission, submit_candidate # noqa: E402
|
||||
from modelhub_client import ModelHubAPIError, ModelHubClientPool, is_duplicate_submission_error # noqa: E402
|
||||
from main import build_adaptive_scan_stages, build_parser, make_run_dir, run_submission, submit_candidate # noqa: E402
|
||||
from modelhub_client import ( # noqa: E402
|
||||
ModelHubAPIError,
|
||||
ModelHubClientPool,
|
||||
OldModelQueuePolicyError,
|
||||
is_duplicate_submission_error,
|
||||
)
|
||||
from models import HFModelSummary, ModelInspection # noqa: E402
|
||||
from outcome_tracker import OutcomeTracker # noqa: E402
|
||||
from submission_claims import SubmissionClaimStore, candidate_key # noqa: E402
|
||||
@@ -210,6 +215,109 @@ def make_candidate(index: int) -> dict:
|
||||
|
||||
|
||||
class ClientPoolConcurrencyTests(unittest.TestCase):
|
||||
def test_old_models_use_only_each_accounts_first_eighty_queue_positions(self) -> None:
|
||||
below_threshold = FakeClient(active_count=79)
|
||||
at_threshold = FakeClient(active_count=80)
|
||||
pool = ModelHubClientPool(
|
||||
[below_threshold, at_threshold], # type: ignore[list-item]
|
||||
active_task_cap=100,
|
||||
active_counts_ttl=60,
|
||||
old_model_queue_threshold=80,
|
||||
recent_model_days=7,
|
||||
instance_id="old-model-threshold-test",
|
||||
)
|
||||
submitted_at = datetime(2026, 8, 11, tzinfo=timezone.utc)
|
||||
|
||||
pool.add_task_for_model(
|
||||
{"model": "old-allowed-as-position-80"},
|
||||
model_last_modified=datetime(2026, 7, 1, tzinfo=timezone.utc),
|
||||
submitted_at=submitted_at,
|
||||
)
|
||||
|
||||
self.assertEqual(1, len(below_threshold.submitted))
|
||||
self.assertEqual(0, len(at_threshold.submitted))
|
||||
self.assertEqual(0, pool.old_model_submit_slots())
|
||||
with self.assertRaises(OldModelQueuePolicyError):
|
||||
pool.add_task_for_model(
|
||||
{"model": "old-rejected"},
|
||||
model_last_modified=datetime(2026, 7, 1, tzinfo=timezone.utc),
|
||||
submitted_at=submitted_at,
|
||||
)
|
||||
|
||||
pool.add_task_for_model(
|
||||
{"model": "recent-allowed"},
|
||||
model_last_modified=datetime(2026, 8, 10, tzinfo=timezone.utc),
|
||||
submitted_at=submitted_at,
|
||||
)
|
||||
self.assertEqual(2, len(below_threshold.submitted) + len(at_threshold.submitted))
|
||||
|
||||
def test_concurrent_old_model_submissions_cannot_cross_eighty(self) -> None:
|
||||
client = FakeClient(active_count=78)
|
||||
pool = ModelHubClientPool(
|
||||
[client], # type: ignore[list-item]
|
||||
active_task_cap=100,
|
||||
active_counts_ttl=60,
|
||||
old_model_queue_threshold=80,
|
||||
recent_model_days=7,
|
||||
)
|
||||
submitted_at = datetime(2026, 8, 11, tzinfo=timezone.utc)
|
||||
|
||||
def submit(index: int) -> str:
|
||||
try:
|
||||
pool.add_task_for_model(
|
||||
{"model": f"old-{index}"},
|
||||
model_last_modified=datetime(2026, 7, 1, tzinfo=timezone.utc),
|
||||
submitted_at=submitted_at,
|
||||
)
|
||||
return "submitted"
|
||||
except OldModelQueuePolicyError:
|
||||
return "policy_skipped"
|
||||
|
||||
with ThreadPoolExecutor(max_workers=8) as executor:
|
||||
outcomes = list(executor.map(submit, range(8)))
|
||||
|
||||
self.assertEqual(2, outcomes.count("submitted"))
|
||||
self.assertEqual(6, outcomes.count("policy_skipped"))
|
||||
self.assertEqual(2, len(client.submitted))
|
||||
self.assertEqual(0, pool.old_model_submit_slots())
|
||||
|
||||
def test_scan_does_not_expand_beyond_recent_window_when_old_slots_are_full(self) -> None:
|
||||
now = datetime(2026, 8, 11, tzinfo=timezone.utc)
|
||||
stages = build_adaptive_scan_stages(
|
||||
now=now,
|
||||
initial_updated_after=now - timedelta(hours=48),
|
||||
initial_limit=100,
|
||||
allow_older_than_recent_window=False,
|
||||
recent_model_days=7,
|
||||
)
|
||||
|
||||
self.assertEqual(["configured_window", "last_7_days"], [stage["name"] for stage in stages])
|
||||
self.assertTrue(all(stage["updatedAfter"] >= now - timedelta(days=7) for stage in stages))
|
||||
|
||||
def test_submit_candidate_reports_old_model_policy_skip_at_eighty(self) -> None:
|
||||
client = FakeClient(active_count=80)
|
||||
pool = ModelHubClientPool(
|
||||
[client], # type: ignore[list-item]
|
||||
active_task_cap=100,
|
||||
old_model_queue_threshold=80,
|
||||
recent_model_days=7,
|
||||
)
|
||||
result = submit_candidate(
|
||||
{
|
||||
"repoId": "owner/old",
|
||||
"modelAddress": "https://modelscope.cn/models/owner/old",
|
||||
"targetGpu": "Iluvatar_bi-150",
|
||||
"taskType": "text-generation",
|
||||
"framework": "vllm",
|
||||
"configParams": "framework: vllm",
|
||||
"lastModified": "2025-01-01T00:00:00+00:00",
|
||||
},
|
||||
pool, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
self.assertEqual("old_model_policy_skipped", result["outcome"])
|
||||
self.assertEqual([], client.submitted)
|
||||
|
||||
def test_online_submission_does_not_construct_llm_even_when_key_is_present(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
root = Path(temporary_dir)
|
||||
|
||||
@@ -47,7 +47,27 @@ class RateLimitedPageClient:
|
||||
return {"success": True, "data": {"models": items}}
|
||||
|
||||
|
||||
class ModelMetadataClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
def request_json(self, _method: str, _path: str, **_kwargs) -> dict: # noqa: ANN003
|
||||
self.calls += 1
|
||||
return {"Code": 200, "Data": {"LastUpdatedTime": 1786294389}}
|
||||
|
||||
|
||||
class ModelScopeDiscoveryTests(unittest.TestCase):
|
||||
def test_model_last_modified_uses_detail_api_and_is_cached(self) -> None:
|
||||
metadata_client = ModelMetadataClient()
|
||||
discovery = HuggingFaceDiscovery(legacy_http_client=metadata_client) # type: ignore[arg-type]
|
||||
|
||||
first = discovery.get_model_last_modified("owner/model")
|
||||
second = discovery.get_model_last_modified("owner/model")
|
||||
|
||||
self.assertEqual(first, second)
|
||||
self.assertEqual(1786294389, int(first.timestamp())) # type: ignore[union-attr]
|
||||
self.assertEqual(1, metadata_client.calls)
|
||||
|
||||
def test_openapi_page_size_never_exceeds_platform_limit(self) -> None:
|
||||
http_client = RecordingHttpClient()
|
||||
discovery = HuggingFaceDiscovery(http_client=http_client) # type: ignore[arg-type]
|
||||
|
||||
46
tests/test_poll_policy.py
Normal file
46
tests/test_poll_policy.py
Normal file
@@ -0,0 +1,46 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
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 poll_runner import resolve_age_cleanup_policy # noqa: E402
|
||||
|
||||
|
||||
class PollPolicyTests(unittest.TestCase):
|
||||
def test_age_cleanup_uses_eighty_once_then_ninety_five(self) -> None:
|
||||
args = argparse.Namespace(
|
||||
old_model_queue_threshold=80,
|
||||
dynamic_old_model_cleanup_threshold=95,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
("initial", 80),
|
||||
resolve_age_cleanup_policy(args, initial_cleanup_pending=True),
|
||||
)
|
||||
self.assertEqual(
|
||||
("dynamic", 95),
|
||||
resolve_age_cleanup_policy(args, initial_cleanup_pending=False),
|
||||
)
|
||||
|
||||
def test_dynamic_cleanup_cannot_be_stricter_than_admission(self) -> None:
|
||||
args = argparse.Namespace(
|
||||
old_model_queue_threshold=80,
|
||||
dynamic_old_model_cleanup_threshold=70,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
("dynamic", 80),
|
||||
resolve_age_cleanup_policy(args, initial_cleanup_pending=False),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -12,16 +13,28 @@ if str(PACKAGE_DIR) in sys.path:
|
||||
sys.path.insert(0, str(PACKAGE_DIR))
|
||||
|
||||
from modelhub_client import ModelHubClient, ModelHubClientPool # noqa: E402
|
||||
from queue_cleanup import OwnedTask, cleanup_certain_oom_tasks, find_certain_oom_tasks # noqa: E402
|
||||
from queue_cleanup import ( # noqa: E402
|
||||
OwnedTask,
|
||||
cleanup_certain_oom_tasks,
|
||||
find_certain_oom_tasks,
|
||||
find_old_overflow_tasks,
|
||||
)
|
||||
|
||||
|
||||
GIB = 1024**3
|
||||
|
||||
|
||||
class FakeQueueClient:
|
||||
def __init__(self, records: list[dict[str, Any]], *, disappear_on_recheck: bool = False) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
records: list[dict[str, Any]],
|
||||
*,
|
||||
disappear_on_recheck: bool = False,
|
||||
drop_first_on_recheck: bool = False,
|
||||
) -> None:
|
||||
self.records = list(records)
|
||||
self.disappear_on_recheck = disappear_on_recheck
|
||||
self.drop_first_on_recheck = drop_first_on_recheck
|
||||
self.waiting_reads = 0
|
||||
self.stopped: list[list[int]] = []
|
||||
|
||||
@@ -32,6 +45,8 @@ class FakeQueueClient:
|
||||
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}}
|
||||
@@ -47,8 +62,13 @@ class FakeQueueClient:
|
||||
|
||||
|
||||
class FakeDiscovery:
|
||||
def __init__(self, sizes: dict[str, int | None]) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
sizes: dict[str, int | None],
|
||||
last_modified: dict[str, datetime | None] | None = None,
|
||||
) -> None:
|
||||
self.sizes = sizes
|
||||
self.last_modified = last_modified or {}
|
||||
|
||||
def list_repo_tree(self, repo_id: str) -> list[dict[str, Any]]:
|
||||
size = self.sizes[repo_id]
|
||||
@@ -56,6 +76,9 @@ class FakeDiscovery:
|
||||
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)
|
||||
|
||||
|
||||
class RecordingHttpClient:
|
||||
def __init__(self) -> None:
|
||||
@@ -74,6 +97,151 @@ class RecordingHttpClient:
|
||||
|
||||
|
||||
class QueueCleanupTests(unittest.TestCase):
|
||||
def test_old_models_are_selected_only_after_each_accounts_first_eighty_tasks(self) -> None:
|
||||
now = datetime(2026, 8, 11, tzinfo=timezone.utc)
|
||||
tasks = [
|
||||
OwnedTask(0, index, "owner/old", "Iluvatar_bi-100", "waiting")
|
||||
for index in range(1, 82)
|
||||
]
|
||||
tasks.append(OwnedTask(0, 82, "owner/recent", "Iluvatar_bi-100", "waiting"))
|
||||
|
||||
selected, skipped = find_old_overflow_tasks(
|
||||
tasks,
|
||||
model_last_modified={
|
||||
"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc),
|
||||
"owner/recent": datetime(2026, 8, 10, tzinfo=timezone.utc),
|
||||
},
|
||||
queue_threshold=80,
|
||||
recent_model_days=7,
|
||||
reference_time=now,
|
||||
)
|
||||
|
||||
self.assertEqual([81], [item["taskId"] for item in selected])
|
||||
self.assertEqual(81, selected[0]["queuePosition"])
|
||||
self.assertEqual(1, skipped["recentOverflowTasks"])
|
||||
|
||||
def test_old_overflow_task_is_not_stopped_if_it_moves_into_first_eighty(self) -> None:
|
||||
records = [
|
||||
{
|
||||
"taskId": index,
|
||||
"modelId": "owner/old",
|
||||
"gpuType": "Iluvatar_bi-100",
|
||||
"status": "waiting",
|
||||
}
|
||||
for index in range(1, 82)
|
||||
]
|
||||
client = FakeQueueClient(records, drop_first_on_recheck=True)
|
||||
pool = ModelHubClientPool(
|
||||
[client], # type: ignore[list-item]
|
||||
active_task_cap=100,
|
||||
old_model_queue_threshold=80,
|
||||
)
|
||||
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(1, summary["oldOverflowCount"])
|
||||
self.assertEqual(0, summary["cancelledCount"])
|
||||
self.assertEqual(1, summary["policyNoLongerAppliesCount"])
|
||||
self.assertEqual([], client.stopped)
|
||||
|
||||
def test_cleanup_stops_old_task_beyond_eightieth_position(self) -> None:
|
||||
records = [
|
||||
{
|
||||
"taskId": index,
|
||||
"modelId": "owner/old",
|
||||
"gpuType": "Iluvatar_bi-100",
|
||||
"status": "waiting",
|
||||
}
|
||||
for index in range(1, 82)
|
||||
]
|
||||
client = FakeQueueClient(records)
|
||||
pool = ModelHubClientPool(
|
||||
[client], # type: ignore[list-item]
|
||||
active_task_cap=100,
|
||||
old_model_queue_threshold=80,
|
||||
)
|
||||
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(1, summary["oldOverflowCount"])
|
||||
self.assertEqual(1, summary["cancelledCount"])
|
||||
self.assertEqual([[81]], client.stopped)
|
||||
|
||||
def test_dynamic_cleanup_keeps_positions_through_ninety_five(self) -> None:
|
||||
records = [
|
||||
{
|
||||
"taskId": index,
|
||||
"modelId": "owner/old",
|
||||
"gpuType": "Iluvatar_bi-100",
|
||||
"status": "waiting",
|
||||
}
|
||||
for index in range(1, 97)
|
||||
]
|
||||
client = FakeQueueClient(records)
|
||||
pool = ModelHubClientPool(
|
||||
[client], # type: ignore[list-item]
|
||||
active_task_cap=100,
|
||||
old_model_queue_threshold=80,
|
||||
)
|
||||
summary = cleanup_certain_oom_tasks(
|
||||
pool,
|
||||
FakeDiscovery(
|
||||
{"owner/old": 1 * GIB},
|
||||
{"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
|
||||
), # type: ignore[arg-type]
|
||||
age_queue_threshold=95,
|
||||
reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
|
||||
log=lambda _message: None,
|
||||
)
|
||||
|
||||
self.assertEqual(95, summary["oldModelQueueThreshold"])
|
||||
self.assertEqual([96], [item["taskId"] for item in summary["oldOverflowTasks"]])
|
||||
self.assertEqual([[96]], client.stopped)
|
||||
|
||||
def test_oom_is_removed_before_recalculating_old_overflow_positions(self) -> None:
|
||||
records = [
|
||||
{
|
||||
"taskId": index,
|
||||
"modelId": "owner/large" if index == 1 else "owner/old",
|
||||
"gpuType": "Iluvatar_bi-100",
|
||||
"status": "waiting",
|
||||
}
|
||||
for index in range(1, 83)
|
||||
]
|
||||
client = FakeQueueClient(records)
|
||||
pool = ModelHubClientPool(
|
||||
[client], # type: ignore[list-item]
|
||||
active_task_cap=100,
|
||||
old_model_queue_threshold=80,
|
||||
)
|
||||
summary = cleanup_certain_oom_tasks(
|
||||
pool,
|
||||
FakeDiscovery(
|
||||
{"owner/large": 40 * GIB, "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(1, summary["certainOomCount"])
|
||||
self.assertEqual([82], [item["taskId"] for item in summary["oldOverflowTasks"]])
|
||||
self.assertEqual([[1], [82]], 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]
|
||||
|
||||
Reference in New Issue
Block a user