177 lines
7.1 KiB
Python
177 lines
7.1 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import datetime, timezone
|
|
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 common import read_jsonl # noqa: E402
|
|
from main import make_run_dir # noqa: E402
|
|
from modelhub_client import ModelHubAPIError, ModelHubClientPool # noqa: E402
|
|
from outcome_tracker import OutcomeTracker # noqa: E402
|
|
from submission_claims import SubmissionClaimStore, candidate_key # noqa: E402
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self, active_count: int = 0, *, reject_capacity: bool = False) -> None:
|
|
self.active_count = active_count
|
|
self.reject_capacity = reject_capacity
|
|
self.submitted: list[dict] = []
|
|
self.count_calls = 0
|
|
self.search_calls = 0
|
|
self._lock = threading.Lock()
|
|
|
|
def count_active_tasks(self, **_kwargs) -> int: # noqa: ANN003
|
|
with self._lock:
|
|
self.count_calls += 1
|
|
return self.active_count
|
|
|
|
def add_task(self, payload: dict) -> dict:
|
|
if self.reject_capacity:
|
|
raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限")
|
|
with self._lock:
|
|
self.submitted.append(payload)
|
|
return {"code": 0, "data": {"id": len(self.submitted)}}
|
|
|
|
def search_by_model_id(self, _model_id: str) -> dict:
|
|
with self._lock:
|
|
self.search_calls += 1
|
|
return {"code": 0, "data": {"verifyResult": {}}}
|
|
|
|
def list_tasks(self, **_kwargs) -> list[dict]: # noqa: ANN003
|
|
return []
|
|
|
|
def list_tasks_page(self, **_kwargs) -> dict: # noqa: ANN003
|
|
return {"code": 0, "data": {"records": [], "pages": 0}}
|
|
|
|
def find_recent_task_id(self, *_args, **_kwargs): # noqa: ANN002, ANN003
|
|
return None
|
|
|
|
|
|
def make_candidate(index: int) -> dict:
|
|
return {
|
|
"repoId": f"owner/model-{index}",
|
|
"modelAddress": f"https://modelscope.cn/models/owner/model-{index}",
|
|
"targetGpu": "NVIDIA-A100",
|
|
}
|
|
|
|
|
|
class ClientPoolConcurrencyTests(unittest.TestCase):
|
|
def test_every_account_reaches_capacity_under_load(self) -> None:
|
|
clients = [FakeClient() for _ in range(12)]
|
|
pool = ModelHubClientPool(
|
|
clients, # type: ignore[arg-type]
|
|
active_task_cap=10,
|
|
active_counts_ttl=60,
|
|
instance_id="full-load-test",
|
|
)
|
|
|
|
with ThreadPoolExecutor(max_workers=24) as executor:
|
|
list(executor.map(lambda index: pool.add_task({"index": index}), range(120)))
|
|
|
|
self.assertEqual([10] * 12, [len(client.submitted) for client in clients])
|
|
self.assertEqual([10] * 12, pool.active_task_counts())
|
|
self.assertEqual(0, pool.available_submit_slots())
|
|
|
|
def test_concurrent_submissions_reserve_and_balance_accounts(self) -> None:
|
|
clients = [FakeClient(), FakeClient()]
|
|
pool = ModelHubClientPool(
|
|
clients, # type: ignore[arg-type]
|
|
active_task_cap=4,
|
|
active_counts_ttl=60,
|
|
instance_id="balance-test",
|
|
)
|
|
|
|
with ThreadPoolExecutor(max_workers=8) as executor:
|
|
results = list(executor.map(lambda index: pool.add_task({"index": index}), range(8)))
|
|
|
|
self.assertEqual(8, len(results))
|
|
self.assertEqual([4, 4], [len(client.submitted) for client in clients])
|
|
self.assertEqual([4, 4], pool.active_task_counts())
|
|
self.assertEqual(0, pool.available_submit_slots())
|
|
# A lagging remote count must not erase successful local reservations.
|
|
pool._refresh_active_counts(force=True)
|
|
self.assertEqual([4, 4], pool.active_task_counts())
|
|
|
|
def test_capacity_rejection_falls_through_to_another_account(self) -> None:
|
|
full_elsewhere = FakeClient(active_count=0, reject_capacity=True)
|
|
available = FakeClient(active_count=1)
|
|
pool = ModelHubClientPool(
|
|
[full_elsewhere, available], # type: ignore[arg-type]
|
|
active_task_cap=2,
|
|
instance_id="fallback-test",
|
|
)
|
|
|
|
response = pool.add_task({"model": "x"})
|
|
|
|
self.assertEqual(1, response["data"]["id"])
|
|
self.assertEqual(0, len(full_elsewhere.submitted))
|
|
self.assertEqual(1, len(available.submitted))
|
|
|
|
def test_verification_cache_is_reset_between_cycles(self) -> None:
|
|
client = FakeClient()
|
|
pool = ModelHubClientPool([client], active_task_cap=2) # type: ignore[arg-type]
|
|
pool.search_by_model_id("owner/model")
|
|
pool.search_by_model_id("owner/model")
|
|
self.assertEqual(1, client.search_calls)
|
|
pool.begin_cycle()
|
|
pool.search_by_model_id("owner/model")
|
|
self.assertEqual(2, client.search_calls)
|
|
|
|
|
|
class ProcessCoordinationTests(unittest.TestCase):
|
|
def test_concurrent_claim_stores_select_disjoint_candidates(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
path = Path(temporary_dir) / "claims.jsonl"
|
|
candidates = [make_candidate(index) for index in range(6)]
|
|
first = SubmissionClaimStore(path, owner_id="worker-a")
|
|
second = SubmissionClaimStore(path, owner_id="worker-b")
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
|
first_future = executor.submit(first.claim, candidates, limit=3)
|
|
second_future = executor.submit(second.claim, candidates, limit=3)
|
|
first_claims = first_future.result()
|
|
second_claims = second_future.result()
|
|
|
|
first_keys = {candidate_key(candidate) for candidate in first_claims}
|
|
second_keys = {candidate_key(candidate) for candidate in second_claims}
|
|
self.assertEqual(3, len(first_claims))
|
|
self.assertEqual(3, len(second_claims))
|
|
self.assertTrue(first_keys.isdisjoint(second_keys))
|
|
|
|
def test_outcome_saves_merge_instead_of_overwriting_other_worker(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
path = Path(temporary_dir) / "outcomes.jsonl"
|
|
first = OutcomeTracker(path)
|
|
second = OutcomeTracker(path)
|
|
first.record_submission("model-a", "gpu", "vllm", "text-generation", "task-a", "2026-01-01T00:00:00+00:00")
|
|
second.record_submission("model-b", "gpu", "vllm", "text-generation", "task-b", "2026-01-01T00:00:01+00:00")
|
|
|
|
first.save()
|
|
second.save()
|
|
|
|
self.assertEqual({"task-a", "task-b"}, {row["taskId"] for row in read_jsonl(path)})
|
|
|
|
def test_run_directories_are_allocated_atomically(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
base = Path(temporary_dir) / "runs"
|
|
now = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
|
with ThreadPoolExecutor(max_workers=8) as executor:
|
|
paths = list(executor.map(lambda _index: make_run_dir(base, now), range(8)))
|
|
|
|
self.assertEqual(8, len(set(paths)))
|
|
self.assertTrue(all(path.is_dir() for path in paths))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|