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()