fix: coordinate concurrent account capacity filling

tests/test_concurrency.py

@@ -0,0 +1,176 @@
+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()